feat: Fix CPU offloading + add options for FSDP offload and activation ckpting

.github/workflows/cicd-main.yml

@@ -150,7 +150,7 @@ jobs:
     if: ${{ needs.pre-flight.outputs.test_level != 'none' }}
     with:
       RUNNER: self-hosted-azure
-      TIMEOUT: 20
+      TIMEOUT: 30
       UNIT_TEST_SCRIPT: |
         cd /opt/reinforcer
         if [[ "${{ needs.pre-flight.outputs.test_level }}" =~ ^(L0|L1|L2)$ ]]; then

examples/configs/grpo_math_1B.yaml

@@ -33,6 +33,8 @@ policy:
   logprob_batch_size: 4
   max_total_sequence_length: 512
   precision: "bfloat16"
+  fsdp_offload_enabled: false
+  activation_checkpointing_enabled: false
 
   optimizer:
     name: "torch.optim.AdamW"

examples/configs/grpo_math_8B.yaml

@@ -15,6 +15,8 @@ policy:
   logprob_batch_size: 2
   max_total_sequence_length: 4096
   precision: "bfloat16"
+  fsdp_offload_enabled: false
+  activation_checkpointing_enabled: false
 
   optimizer:
     name: "torch.optim.AdamW"

examples/configs/sft.yaml

@@ -25,6 +25,8 @@ policy:
   train_micro_batch_size: 1
   max_total_sequence_length: 1024
   precision: "float32"
+  fsdp_offload_enabled: false
+  activation_checkpointing_enabled: false
 
   optimizer:
     name: "torch.optim.AdamW"

nemo_reinforcer/distributed/worker_groups.py

@@ -135,7 +135,8 @@ def __call__(
             if env_vars:
                 if "runtime_env" not in options:
                     options["runtime_env"] = {}
-                options["runtime_env"]["env_vars"] = env_vars
+                for k, v in env_vars.items():
+                    options["runtime_env"]["env_vars"][k] = v
 
             # Apply initialization parameters
             if init_kwargs:
@@ -207,7 +208,6 @@ def __init__(
         self.cluster = cluster
         self.name_prefix = name_prefix
         self.tied_workers_groups = []
-
         # Maps worker indices to their corresponding tied group index
         # For example, if worker with index 3 belongs to tied worker group 1,
         # then worker_to_tied_group_index[3] = 1

nemo_reinforcer/models/policy/__init__.py

@@ -31,3 +31,5 @@ class PolicyConfig(TypedDict):
     logprob_batch_size: int
     generation: GenerationConfig
     precision: str
+    fsdp_offload_enabled: bool
+    activation_checkpointing_enabled: bool

nemo_reinforcer/models/policy/hf_policy.py

@@ -22,6 +22,7 @@
 import torch
 from torch.distributed.device_mesh import init_device_mesh
 from torch.distributed.fsdp import (
+    CPUOffload,
     FullyShardedDataParallel,
     MixedPrecision,
 )
@@ -110,6 +111,12 @@ def __init__(
         # ------------------------------------------------
 
         def do_fsdp(model):
+            if world_size == 1:
+                print(
+                    "[INFO] Using a single GPU - skipping FSDP wrapper to avoid GPU memory offloading issues"
+                )
+                return model
+
             # Create a device mesh with 'world_size' GPUs in a 1D arrangement.
             mesh = init_device_mesh("cuda", (world_size,))
             mp_policy = MixedPrecision(
@@ -118,21 +125,32 @@ def do_fsdp(model):
                 buffer_dtype=torch.float32,
             )
 
+            cpu_offload = (
+                CPUOffload(offload_params=True)
+                if self.cfg["fsdp_offload_enabled"]
+                else None
+            )
+
             return FullyShardedDataParallel(
                 model,
                 device_mesh=mesh,
                 auto_wrap_policy=size_based_auto_wrap_policy,
                 mixed_precision=mp_policy,
+                cpu_offload=cpu_offload,
             )
 
         self.model.to("cuda")
+        if self.cfg["activation_checkpointing_enabled"]:
+            self.model.gradient_checkpointing_enable(
+                gradient_checkpointing_kwargs={"use_reentrant": False}
+            )
         self.model = do_fsdp(self.model)
-        self.model = self.move_to_cpu(self.model)
+        self.model = self.manual_offload_to_cpu(self.model)
         if self.reference_model is not None:
             self.reference_model.to("cuda")
             self.reference_model = do_fsdp(self.reference_model)
-            self.reference_model = self.move_to_cpu(self.reference_model)
-        self.model.to("cuda")
+            self.reference_model = self.manual_offload_to_cpu(self.reference_model)
+        self.model = self.manual_load_to_gpu(self.model)
         self._held_reference_model_params = None
         # register_fsdp_forward_method(self.model, "generate")
         if init_optimizer:
@@ -440,8 +458,8 @@ def use_reference_model(self):
             original_model = self.model
             original_reference_model = self.reference_model
 
-            self.model = self.move_to_cpu(self.model)
-            self.reference_model = self.reference_model.to("cuda")
+            self.model = self.manual_offload_to_cpu(self.model)
+            self.reference_model = self.manual_load_to_gpu(self.reference_model)
 
             # Swap the references
             self.model, self.reference_model = self.reference_model, self.model
@@ -454,8 +472,8 @@ def use_reference_model(self):
 
         finally:
             # Restore original references and device placement
-            self.reference_model = self.move_to_cpu(original_reference_model)
-            self.model = original_model.to("cuda")
+            self.reference_model = self.manual_offload_to_cpu(original_reference_model)
+            self.model = self.manual_load_to_gpu(original_model)
             gc.collect()
             torch.cuda.empty_cache()
 
@@ -537,7 +555,13 @@ def generate(
                     # Set attention mask for the actual tokens (at the end for left padding)
                     left_padded_attention_mask[i, seq_len - length :] = 1
 
-                outputs = self.model.module.generate(
+                if isinstance(
+                    self.model, torch.distributed.fsdp.FullyShardedDataParallel
+                ):
+                    generation_module = self.model.module
+                else:
+                    generation_module = self.model
+                outputs = generation_module.generate(
                     input_ids=left_padded_input_ids,
                     attention_mask=left_padded_attention_mask,
                     max_new_tokens=gen_cfg["max_new_tokens"],
@@ -724,6 +748,11 @@ def report_device_id(self) -> str:
     def get_weight_ipc_handles(self, offload_model=True):
         from torch.multiprocessing.reductions import reduce_tensor
 
+        # If the model is not FSDP, then we need to manually move it to the GPU
+        # For an FSDP model, model.state_dict() will move the params to the GPU
+        if not isinstance(self.model, torch.distributed.fsdp.FullyShardedDataParallel):
+            self.model = self.manual_load_to_gpu(self.model)
+
         # TODO @sahilj: do this without an allgather (maybe FSDP2)
         params = self.model.state_dict()
 
@@ -735,49 +764,50 @@ def get_weight_ipc_handles(self, offload_model=True):
 
         # Replace the original params with the converted ones
         params = dtype_params
+        # For FSDP1, params may get GC'ed before sending to vllm,
+        # so we need to hold a reference to them
         self._held_reference_model_params = params
         data = {}
         device_uuid = self.report_device_id()
         for name, p in params.items():
             data[name] = reduce_tensor(p.detach())
 
         if offload_model:
-            self.model = self.move_to_cpu(self.model)
+            self.model = self.manual_offload_to_cpu(self.model)
             gc.collect()
             torch.cuda.empty_cache()
         return {device_uuid: data}
 
     def prepare_for_lp_inference(self):
-        self.model.to("cuda")
+        self.model = self.manual_load_to_gpu(self.model)
         self.model.eval()
         self.offload_before_refit()
 
     def prepare_for_training(self, *args, **kwargs):
         # onload models and optimizer state to cuda
-        self.model.to("cuda")
+        self.model = self.manual_load_to_gpu(self.model)
         self.model.train()
 
-        # Move optimizer state to CUDA if it exists
-        if hasattr(self, "optimizer") and self.optimizer is not None:
-            for state in self.optimizer.state.values():
-                for k, v in state.items():
-                    if torch.is_tensor(v) and not v.is_cuda:
-                        state[k] = v.to("cuda")
+        if not self.cfg["fsdp_offload_enabled"]:
+            # Move optimizer state to CUDA if it exists
+            if hasattr(self, "optimizer") and self.optimizer is not None:
+                for state in self.optimizer.state.values():
+                    for k, v in state.items():
+                        if torch.is_tensor(v) and not v.is_cuda:
+                            state[k] = v.to("cuda")
 
         torch.cuda.empty_cache()
 
     @torch.no_grad()
     def offload_before_refit(self):
         """Offload the optimizer and buffers to the CPU."""
         torch.randn(1).cuda()  # wake up torch allocator
-        if hasattr(self, "optimizer") and self.optimizer is not None:
-            for state in self.optimizer.state.values():
-                for k, v in state.items():
-                    if torch.is_tensor(v):
-                        state[k] = v.to("cpu")
-
-        for buffer in self.model.buffers():
-            buffer.data = buffer.data.to("cpu")
+        if not self.cfg["fsdp_offload_enabled"]:
+            if hasattr(self, "optimizer") and self.optimizer is not None:
+                for state in self.optimizer.state.values():
+                    for k, v in state.items():
+                        if torch.is_tensor(v):
+                            state[k] = v.to("cpu")
 
         gc.collect()
         torch.cuda.empty_cache()
@@ -792,7 +822,7 @@ def offload_before_refit(self):
     @torch.no_grad()
     def offload_after_refit(self):
         # Offload as much as possible on the CPU
-        self.model = self.move_to_cpu(self.model)
+        self.model = self.manual_offload_to_cpu(self.model)
         self.model.eval()
         torch.randn(1).cuda()  # wake up torch allocator
         self.offload_before_refit()  # rerun the old offload function
@@ -810,15 +840,39 @@ def offload_after_refit(self):
             f"GPU Memory after refit complete: {allocated:.2f}GB allocated, {reserved:.2f}GB reserved"
         )
 
-    def move_to_cpu(self, model):
+    def manual_offload_to_cpu(self, model):
+        if self.cfg["fsdp_offload_enabled"]:
+            return model
+
         for param in model.parameters():
-            param.data = param.data.to("cpu")
+            param.data = param.data.to("cpu", non_blocking=True)
+            if hasattr(param, "_local_shard"):
+                param._local_shard = param.data
+            if param.grad is not None:
+                param.grad = param.grad.to("cpu", non_blocking=True)
+        for buffer in model.buffers():
+            buffer.data = buffer.data.to("cpu", non_blocking=True)
+
+        if hasattr(model, "_fsdp_wrapped_module"):
+            self.manual_offload_to_cpu(model._fsdp_wrapped_module)
 
+        return model
+
+    def manual_load_to_gpu(self, model):
+        if self.cfg["fsdp_offload_enabled"]:
+            return model
+
+        for param in model.parameters():
+            param.data = param.data.to("cuda", non_blocking=True)
+            if hasattr(param, "_local_shard"):
+                param._local_shard = param.data
+            if param.grad is not None:
+                param.grad = param.grad.to("cuda", non_blocking=True)
         for buffer in model.buffers():
-            buffer.data = buffer.data.to("cpu")
+            buffer.data = buffer.data.to("cuda", non_blocking=True)
 
         if hasattr(model, "_fsdp_wrapped_module"):
-            model._fsdp_wrapped_module.to("cpu")
+            self.manual_load_to_gpu(model._fsdp_wrapped_module)
 
         return model
 

nemo_reinforcer/utils/native_checkpoint.py

@@ -154,7 +154,10 @@ def save_checkpoint(
         save_hf: Whether to save in HuggingFace format
     """
     if save_hf:
-        model_state_dict = model._fsdp_wrapped_module.state_dict()
+        if hasattr(model, "_fsdp_wrapped_module"):
+            model_state_dict = model._fsdp_wrapped_module.state_dict()
+        else:
+            model_state_dict = model.state_dict()
 
         if torch.distributed.get_rank() == 0:
             # Create a new path by appending "-hf" to the weights path

tests/unit/models/generation/test_vllm_generation.py

@@ -61,6 +61,8 @@
     "max_new_tokens": 16,
     "do_sample": False,
     "precision": "float32",
+    "fsdp_offload_enabled": False,
+    "activation_checkpointing_enabled": False,
     "optimizer": {
         "name": "torch.optim.AdamW",
         "kwargs": {