feat: Using mcore cpu optimizer (#1242)

Signed-off-by: Guyue Huang <guyueh@nvidia.com>
Signed-off-by: Guyue Huang <140554423+guyueh1@users.noreply.github.com>
Co-authored-by: Terry Kong <terrycurtiskong@gmail.com>

Author: guyueh1

examples/configs/dpo.yaml

@@ -136,6 +136,10 @@ policy:
 
       clip_grad: ${policy.max_grad_norm}
 
+      # optimizer cpu offload
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
     scheduler:
       start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
       end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}

examples/configs/grpo_math_1B.yaml

@@ -110,6 +110,10 @@ policy:
 
       clip_grad: ${policy.max_grad_norm}
 
+      # optimizer cpu offload
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
     scheduler:
       start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
       end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}

examples/configs/rm.yaml

@@ -105,6 +105,10 @@ policy:
 
       clip_grad: ${policy.max_grad_norm}
 
+      # optimizer cpu offload
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
     scheduler:
       start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
       end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}

examples/configs/sft.yaml

@@ -114,6 +114,10 @@ policy:
 
       clip_grad: ${policy.max_grad_norm}
 
+      # optimizer cpu offload
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
     scheduler:
       start_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}
       end_weight_decay: ${policy.megatron_cfg.optimizer.weight_decay}

examples/configs/sft_openmathinstruct2_megatron.yaml

@@ -62,6 +62,10 @@ policy:
       use_precision_aware_optimizer: false #true ## TODO: precision aware optim not working with fp8. Is this expected?
       weight_decay: 0.01
 
+      # optimizer cpu offload
+      optimizer_cpu_offload: false
+      optimizer_offload_fraction: 0.0
+
       ## recently introduced, our current mcore commit doesn't have this
       #fp8_recipe: delayed
 

nemo_rl/models/policy/__init__.py

@@ -61,6 +61,11 @@ class MegatronOptimizerConfig(TypedDict):
     use_distributed_optimizer: bool
     use_precision_aware_optimizer: bool
     clip_grad: float
+    # knob to enable optimizer cpu offload
+    optimizer_cpu_offload: bool
+    # knob to set the fraction of parameters to keep on CPU
+    # currently if optimizer_cpu_offload is true, this knob must be 1.0
+    optimizer_offload_fraction: float
 
 
 class MegatronSchedulerConfig(TypedDict):

nemo_rl/models/policy/megatron_policy_worker.py

@@ -618,6 +618,19 @@ def __init__(
                     "Refer to https://github.com/NVIDIA-NeMo/RL/issues/1164 for latest updates with this issue."
                 )
 
+        optimizer_cpu_offload = self.cfg["megatron_cfg"]["optimizer"][
+            "optimizer_cpu_offload"
+        ]
+        optimizer_offload_fraction = self.cfg["megatron_cfg"]["optimizer"][
+            "optimizer_offload_fraction"
+        ]
+        if optimizer_cpu_offload:
+            # Currently, hybrid optimizer (partly on GPU and partly on CPU) is not supported because it conflicts with the way
+            # Nemo-rl handles the optimizer offload/onload between generation and training. So if using CPU optimizer the offload_fraction should be 1.0.
+            assert optimizer_offload_fraction == 1.0, (
+                "Currently for optimizer offloading, only optimizer_offload_fraction=1.0 is supported"
+            )
+
         checkpoint_config = CheckpointConfig(
             save_interval=100,
             save=weights_path,
@@ -1759,7 +1772,11 @@ def prepare_for_training(self, *args, **kwargs):
         self.model.train()
 
         # Move optimizer state to CUDA if it exists
-        if hasattr(self, "optimizer") and self.optimizer is not None:
+        if (
+            hasattr(self, "optimizer")
+            and self.optimizer is not None
+            and (not self.cfg["megatron_cfg"]["optimizer"]["optimizer_cpu_offload"])
+        ):
             if isinstance(self.optimizer, ChainedOptimizer):
                 optimizer_state = self.optimizer.state
             else:
@@ -1786,7 +1803,11 @@ def offload_before_refit(self):
             self.model, "cpu", move_params=False, move_grads=True
         )  # get rid of grad buffers
         torch.randn(1).cuda()  # wake up torch allocator
-        if hasattr(self, "optimizer") and self.optimizer is not None:
+        if (
+            hasattr(self, "optimizer")
+            and self.optimizer is not None
+            and (not self.cfg["megatron_cfg"]["optimizer"]["optimizer_cpu_offload"])
+        ):
             # Iterate through the state dictionaries for each parameter group
             if isinstance(self.optimizer, ChainedOptimizer):
                 optimizer_state = self.optimizer.state

tests/unit/models/generation/test_vllm_generation.py

@@ -192,6 +192,8 @@ def get_basic_megatron_test_config(
                 "use_distributed_optimizer": True,
                 "use_precision_aware_optimizer": True,
                 "clip_grad": 1.0,
+                "optimizer_cpu_offload": False,
+                "optimizer_offload_fraction": 0.0,
             },
             "scheduler": {
                 "start_weight_decay": 0.01,

tests/unit/models/policy/test_megatron_worker.py

@@ -114,6 +114,8 @@ def create_megatron_test_config(
                 "use_distributed_optimizer": True,
                 "use_precision_aware_optimizer": True,
                 "clip_grad": 1.0,
+                "optimizer_cpu_offload": False,
+                "optimizer_offload_fraction": 0.0,
             },
             "scheduler": {
                 "start_weight_decay": 0.01,

tools/refit_verifier.py

@@ -232,6 +232,9 @@ def setup_configs(args, tokenizer):
                 "use_distributed_optimizer": True,
                 "use_precision_aware_optimizer": True,
                 "clip_grad": 1.0,
+                # Optimizer CPU offload settings
+                "optimizer_cpu_offload": False,
+                "optimizer_offload_fraction": 0.0,
             },
             "scheduler": {
                 "start_weight_decay": 0.01,