wip

Author: kevssim

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

@@ -21,13 +21,14 @@ def __init__(
         mixed_precision: Literal['no', 'fp8', 'fp16', 'bf16'] = 'bf16',
         ddp_config: Dict[str, Any] = None,
         fsdp_config: Dict[str, Any] = None,
+        memory_efficient: bool = True,
     ):
         from accelerate import Accelerator
 
         self.device_mesh = device_mesh
         self.mixed_precision = mixed_precision
         parallelism_config = self._parallelism_config_from_device_mesh(device_mesh)
-        fsdp_plugin = self._fsdp_config_from_device_mesh(device_mesh, fsdp_config)
+        fsdp_plugin = self._fsdp_config_from_device_mesh(device_mesh, fsdp_config, memory_efficient)
 
         kwargs_handlers = []
         if ddp_config is not None:
@@ -69,7 +70,7 @@ def _parallelism_config_from_device_mesh(device_mesh: DeviceMesh):
 
         return parallelism_config
 
-    def _fsdp_config_from_device_mesh(self, device_mesh: DeviceMesh, fsdp_config: Dict[str, Any]):
+    def _fsdp_config_from_device_mesh(self, device_mesh: DeviceMesh, fsdp_config: Dict[str, Any], memory_efficient: bool):
         from accelerate import FullyShardedDataParallelPlugin
         from torch.distributed.fsdp import BackwardPrefetch
         from torch.distributed.fsdp import ShardingStrategy as FSDPShardingStrategy
@@ -107,7 +108,7 @@ def _fsdp_config_from_device_mesh(self, device_mesh: DeviceMesh, fsdp_config: Di
             activation_checkpointing=fsdp_config.pop('activation_checkpointing', False),
             auto_wrap_policy=fsdp_config.pop('auto_wrap_policy', 'transformer_based_wrap'),  # noqa
             reshard_after_forward=fsdp_config.pop('reshard_after_forward', True),
-            cpu_ram_efficient_loading=fsdp_config.pop('cpu_ram_efficient_loading', True),
+            cpu_ram_efficient_loading=fsdp_config.pop('cpu_ram_efficient_loading', memory_efficient),
             **fsdp_config,
         )
         # The env vars (ACCELERATE_USE_FSDP, FSDP_CPU_RAM_EFFICIENT_LOADING) are set

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

@@ -59,6 +59,12 @@ def wrap_model(self, model, optimizer=None, memory_efficient=True):
             if use_meta:
                 original_sd = model.state_dict()
                 saved_buffers = _get_non_persistent_buffers(model)
+                # Drop optimizer references so old params can be freed on to('meta').
+                # Without this, the optimizer holds strong refs to the full-size
+                # parameter tensors, preventing GC even after the model moves to meta.
+                # _rebind_optimizer will re-attach the new sharded params later.
+                if optimizer is not None:
+                    _unbind_optimizer_params(optimizer)
                 model = model.to(torch.device('meta'))
                 if hasattr(model, 'tie_weights'):
                     model.tie_weights()
@@ -506,6 +512,25 @@ def _get_non_persistent_buffers(model: nn.Module) -> Dict[str, torch.Tensor]:
     return {k: v.clone() for k, v in model.named_buffers() if k in non_persistent_fqns}
 
 
+def _unbind_optimizer_params(optimizer: torch.optim.Optimizer) -> None:
+    """Replace optimizer param references with ``torch.empty(1)`` placeholders.
+
+    This drops the optimizer's strong references to the full model parameters,
+    allowing them to be freed when the model is moved to ``meta`` device.
+    Without this, ``model.to('meta')`` cannot free the old parameter tensors
+    because the optimizer still holds references to them.
+
+    Must be called BEFORE ``model.to('meta')``.  After ``fully_shard`` and
+    ``_broadcast_sharded_state_dict``, call ``_rebind_optimizer`` to point
+    the optimizer at the new sharded parameters.
+
+    This mirrors accelerate's approach in ``Accelerator._prepare_fsdp2``.
+    """
+    for group in optimizer.param_groups:
+        for i in range(len(group['params'])):
+            group['params'][i] = torch.empty(1)
+
+
 def _restore_non_persistent_buffers(
     model: nn.Module,
     saved_buffers: Dict[str, torch.Tensor],

src/twinkle/model/transformers/transformers.py

@@ -196,8 +196,8 @@ def __init__(
         self.mixed_precision = mixed_precision
         self._fsdp_config = dict(fsdp_config or {})
         self._ddp_config = ddp_config or {}
-        self._decide_strategy(strategy)
         self._memory_efficient_init = memory_efficient_init
+        self._decide_strategy(strategy)
         self.grad_scaler_config = grad_scaler_config
         if isinstance(model_cls, str):
             model_cls = getattr(transformers, model_cls)
@@ -267,7 +267,8 @@ def _decide_strategy(self, strategy: Literal['accelerate', 'native_fsdp']):
                 mixed_precision=self.mixed_precision,
                 ddp_config=self._ddp_config,
                 fsdp_config=self._fsdp_config,
-                device_mesh=self.device_mesh)
+                device_mesh=self.device_mesh,
+                memory_efficient=self._memory_efficient_init)
 
         # Sequence parallel ("ulysses") is derived from dp/fsdp ranks; it does not change world size.
         # We construct `sp_strategy` after the underlying HF model is initialized (see __init__).