Fix buffer device move when meta load unavailable

Co-authored-by: sami <[email protected]>

Author: cursoragent

src/prime_rl/trainer/model.py

@@ -268,13 +268,6 @@ def load_dcp_from_hf(model: nn.Module, config: ModelConfig, parallel_dims: Paral
     model.to_empty(device=device)
     torch.distributed.barrier()
 
-    def _move_buffers_to_cuda():
-        """FSDP CPU offloading only manages parameters, not buffers. Move buffers to CUDA."""
-        if config.fsdp_cpu_offload:
-            for name, buffer in model.named_buffers():
-                if buffer.device.type == "cpu":
-                    buffer.data = buffer.data.to("cuda")
-
     def _init_buffers_post_meta():
         if isinstance(model, PreTrainedModelPrimeRL):
             model.init_buffers_post_meta()
@@ -285,7 +278,7 @@ def _init_buffers_post_meta():
     if config.debug.random_init:
         logger.warning("Randomly initializing model. Skipping loading weights from HF.")
         _init_buffers_post_meta()
-        _move_buffers_to_cuda()
+        _move_buffers_to_cuda(model, config)
         return
 
     if not Path(config.name).exists():
@@ -347,7 +340,7 @@ def _init_buffers_post_meta():
     )
     _init_buffers_post_meta()
 
-    _move_buffers_to_cuda()
+    _move_buffers_to_cuda(model, config)
 
     lora_modules = [m for m in model.modules() if hasattr(m, "_init_lora_parameters")]
     if lora_modules:
@@ -431,6 +424,15 @@ def apply_ep(model: nn.Module, parallel_dims: ParallelDims):
             )
 
 
+def _move_buffers_to_cuda(model: nn.Module, config: ModelConfig) -> None:
+    """FSDP CPU offloading only manages parameters, not buffers. Move buffers to CUDA."""
+    if not config.fsdp_cpu_offload:
+        return
+    for _, buffer in model.named_buffers():
+        if buffer.device.type == "cpu":
+            buffer.data = buffer.data.to("cuda")
+
+
 def setup_model(
     config: ModelConfig, parallel_dims: ParallelDims, loading_from_checkpoint_later: bool = False
 ) -> nn.Module:
@@ -472,6 +474,9 @@ def setup_model(
 
     setup_fsdp(model, config, parallel_dims)
 
+    if not possible_to_load_to_meta:
+        _move_buffers_to_cuda(model, config)
+
     # 2. if we can load to meta, we either:
     if possible_to_load_to_meta:
         # - load from checkpoint later if needed
@@ -487,12 +492,7 @@ def setup_model(
             else:
                 fix_model_post_empty(model)
 
-            # FSDP CPU offloading only manages parameters, not buffers.
-            # Buffers must be on CUDA for the forward pass.
-            if config.fsdp_cpu_offload:
-                for name, buffer in model.named_buffers():
-                    if buffer.device.type == "cpu":
-                        buffer.data = buffer.data.to("cuda")
+            _move_buffers_to_cuda(model, config)
         # - or load from HF with dcp
         else:
             load_dcp_from_hf(model, config, parallel_dims)