wip

Author: kevssim

src/twinkle/model/transformers/strategy/accelerate.py

@@ -1,8 +1,8 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
-import os
 from typing import Any, Dict, Literal, Optional
 
 from twinkle import DeviceMesh
+from .load_context import fsdp_pretrained_load_context
 
 
 class AccelerateStrategy:
@@ -27,6 +27,7 @@ def __init__(
 
         self.device_mesh = device_mesh
         self.mixed_precision = mixed_precision
+        self._memory_efficient_init = memory_efficient_init
         parallelism_config = self._parallelism_config_from_device_mesh(device_mesh)
         fsdp_plugin = self._fsdp_config_from_device_mesh(device_mesh, fsdp_config, memory_efficient_init)
 
@@ -43,6 +44,9 @@ def __init__(
             kwargs_handlers=kwargs_handlers,
         )
 
+    def pretrained_load_context(self):
+        return fsdp_pretrained_load_context(self._memory_efficient_init and self.device_mesh is not None)
+
     @staticmethod
     def _parallelism_config_from_device_mesh(device_mesh: DeviceMesh):
         # TODO should test with transformers v5.0

src/twinkle/model/transformers/strategy/load_context.py

@@ -0,0 +1,27 @@
+# Copyright (c) ModelScope Contributors. All rights reserved.
+import contextlib
+import os
+
+_FSDP_EFFICIENT_LOADING_ENV = {
+    'ACCELERATE_USE_FSDP': 'true',
+    'FSDP_CPU_RAM_EFFICIENT_LOADING': 'true',
+}
+
+
+@contextlib.contextmanager
+def fsdp_pretrained_load_context(enabled: bool):
+    """Enable the env flags required for transformers FSDP-aware loading when needed."""
+    if not enabled:
+        yield
+        return
+
+    saved_env = {key: os.environ.get(key) for key in _FSDP_EFFICIENT_LOADING_ENV}
+    os.environ.update(_FSDP_EFFICIENT_LOADING_ENV)
+    try:
+        yield
+    finally:
+        for key, old_val in saved_env.items():
+            if old_val is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = old_val

src/twinkle/model/transformers/strategy/native_fsdp.py

@@ -7,6 +7,7 @@
 from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Set
 
 from twinkle.utils import DeviceMesh, Platform, torch_util
+from .load_context import fsdp_pretrained_load_context
 
 if TYPE_CHECKING:
     from torch.distributed.fsdp import MixedPrecisionPolicy
@@ -28,6 +29,9 @@ def __init__(self,
         self.enable_ep = enable_ep
         self.ep_fsdp_device_mesh = self._build_ep_fsdp_device_mesh(ep_size) if enable_ep else None
 
+    def pretrained_load_context(self):
+        return fsdp_pretrained_load_context(self._memory_efficient_init and self.device_mesh is not None)
+
     def _build_ep_fsdp_device_mesh(self, ep_size: Optional[int] = None) -> Optional[TorchDeviceMesh]:
         if self.device_mesh is None:
             return None

src/twinkle/model/transformers/transformers.py

@@ -212,22 +212,8 @@ def __init__(
         else:
             model_id = HubOperation.download_model(model_id)
             # Trigger transformers' FSDP-aware loading: meta-device init + rank-0-only weight load.
-            use_efficient_loading = (memory_efficient_init and self.device_mesh is not None)
-            _saved_env = {}
-            if use_efficient_loading:
-                _saved_env['ACCELERATE_USE_FSDP'] = os.environ.get('ACCELERATE_USE_FSDP')
-                _saved_env['FSDP_CPU_RAM_EFFICIENT_LOADING'] = os.environ.get('FSDP_CPU_RAM_EFFICIENT_LOADING')
-                os.environ['ACCELERATE_USE_FSDP'] = 'true'
-                os.environ['FSDP_CPU_RAM_EFFICIENT_LOADING'] = 'true'
-            try:
+            with self.strategy.pretrained_load_context():
                 self.model = model_cls.from_pretrained(model_id, config=config, **kwargs)
-            finally:
-                if use_efficient_loading:
-                    for key, old_val in _saved_env.items():
-                        if old_val is None:
-                            os.environ.pop(key, None)
-                        else:
-                            os.environ[key] = old_val
         self.model.gradient_checkpointing_enable()
         self.sp_strategy = None
         self._model_wrapped = False

tests/strategy/test_fsdp2_memory_efficient_init.py

@@ -211,7 +211,7 @@ def _worker_wrap_model_memory_efficient(rank, world_size, port, ref_sd):
         mesh_dim_names=('fsdp', ),
         device_type=_DEVICE_TYPE,
     )
-    strategy = NativeFSDPStrategy(device_mesh=mesh, mixed_precision='no', memory_efficient=True)
+    strategy = NativeFSDPStrategy(device_mesh=mesh, mixed_precision='no', memory_efficient_init=True)
 
     model = TinyModel(dim=32).to(_DEVICE_TYPE)
     if rank == 0:
@@ -269,7 +269,7 @@ def _worker_wrap_model_legacy(rank, world_size, port, ref_sd):
         mesh_dim_names=('fsdp', ),
         device_type=_DEVICE_TYPE,
     )
-    strategy = NativeFSDPStrategy(device_mesh=mesh, mixed_precision='no', memory_efficient=False)
+    strategy = NativeFSDPStrategy(device_mesh=mesh, mixed_precision='no', memory_efficient_init=False)
 
     model = TinyModel(dim=32).to(_DEVICE_TYPE)
     model.load_state_dict(ref_sd)
@@ -324,7 +324,7 @@ def _worker_wrap_model_per_layer(rank, world_size, port, ref_sd):
         mesh_dim_names=('fsdp', ),
         device_type=_DEVICE_TYPE,
     )
-    strategy = NativeFSDPStrategy(device_mesh=mesh, mixed_precision='no', memory_efficient=True)
+    strategy = NativeFSDPStrategy(device_mesh=mesh, mixed_precision='no', memory_efficient_init=True)
 
     model = TinyTransformerModel(dim=32, num_layers=2).to(_DEVICE_TYPE)
     if rank == 0: