diff --git a/examples/configs/dpo.yaml b/examples/configs/dpo.yaml index ba4fb18e607..758f263768f 100755 --- a/examples/configs/dpo.yaml +++ b/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} diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 141d33f8e36..494ef7ea4df 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/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} diff --git a/examples/configs/rm.yaml b/examples/configs/rm.yaml index 1a8ab49f29e..7914313fa22 100644 --- a/examples/configs/rm.yaml +++ b/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} diff --git a/examples/configs/sft.yaml b/examples/configs/sft.yaml index 2f6738d6254..ddb77d07f1a 100644 --- a/examples/configs/sft.yaml +++ b/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} diff --git a/examples/configs/sft_openmathinstruct2_megatron.yaml b/examples/configs/sft_openmathinstruct2_megatron.yaml index 67734cebe44..1576c6093b1 100644 --- a/examples/configs/sft_openmathinstruct2_megatron.yaml +++ b/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 diff --git a/nemo_rl/models/policy/__init__.py b/nemo_rl/models/policy/__init__.py index 71f62f95422..994a1d41aee 100644 --- a/nemo_rl/models/policy/__init__.py +++ b/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): diff --git a/nemo_rl/models/policy/megatron_policy_worker.py b/nemo_rl/models/policy/megatron_policy_worker.py index 620d02a41b1..01f12f92039 100644 --- a/nemo_rl/models/policy/megatron_policy_worker.py +++ b/nemo_rl/models/policy/megatron_policy_worker.py @@ -614,6 +614,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, @@ -1753,7 +1766,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: @@ -1780,7 +1797,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 diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index 140efca56bb..2ff6cd01cc4 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/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, diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index e83b5818f2d..2f788846bca 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/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, diff --git a/tools/refit_verifier.py b/tools/refit_verifier.py index f4d1059c84e..475bfd1215d 100644 --- a/tools/refit_verifier.py +++ b/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,