GRPO OOM after checkpointing

[bxyu-nvidia]

Describe the bug

In GRPO, after checkpointing, the subsequent vLLM wakeup experiences an OOM. See https://wandb.ai/nvidia/bxyu-nemo-gym-rl-integration/runs/fbgrwc5h/logs line 3329. There are some unrelated logs there regarding penguin

Steps/Code to reproduce bug

Please list minimal steps or code snippet for us to be able to reproduce the bug.

A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports.

Expected behavior

A clear and concise description of what you expected to happen.

Environment overview (please complete the following information)

• Environment location: [Bare-metal, Docker, Cloud(specify cloud provider - AWS, Azure, GCP, Collab)]
• Method of install: [pip install or from source]. Please specify exact commands you used to install.
• If method of install is [Docker], provide docker pull & docker run commands used

Environment details

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:

• OS version
• PyTorch version
• Python version

Additional context

Add any other context about the problem here.
Example: GPU model

[terrykong]

@bxyu-nvidia to provide a reproducer when you can.

The WAR from @bxyu-nvidia is to lower the vllm gpu memory utilization

[sharonyu-115]

Hi @terrykong @bxyu-nvidia I encountered issue when using the megatron backend. Every time after the checkpoing saving step, vllm wake up got OOM error:

�[36m(MegatronPolicyWorker[rank=2] pid=916738)�[0m [SHARON] GPU Memory before saving checkpoint: 24.89GB allocated, 24.99GB reserved
�[36m(MegatronPolicyWorker[rank=0] pid=916741)�[0m saving checkpoint at iteration       0 to /lustre/fs1/portfolios/coreai/users/shuangy/src/nemo-rl-main/RL/results/Qwen8B-DAPO-Runs_qwen8b-fp8-dapo-deepgemm-cudagraph/tmp_step_50/policy/weights in torch_dist format
�[36m(MegatronPolicyWorker[rank=0] pid=916741)�[0m Storing distributed optimizer sharded state of type fully_sharded_model_space
�[36m(MegatronPolicyWorker[rank=1] pid=916742)�[0m WARNING:megatron.core.dist_checkpointing.validation:Rank 1 common state dict differs from rank 0 common state dict. Keys only on rank 0: [], Keys only on 1: [], Mismatched keys: [(('optimizer', 'optimizer', 'param_groups', 1, 'step'), <class 'int'>, <class 'int'>)]
�[36m(VllmGenerationWorker pid=906502)�[0m 
Processed prompts:  84%|████████▍ | 27/32 [01:21<00:11,  2.25s/it, est. speed input: 48.51 toks/s, output: 1547.18 toks/s]�[32m [repeated 2x across cluster]�[0m
�[36m(VllmGenerationWorker pid=906502)�[0m 
Processed prompts:  97%|█████████▋| 31/32 [01:22<00:01,  1.13s/it, est. speed input: 54.62 toks/s, output: 1932.63 toks/s]
�[36m(MegatronPolicyWorker[rank=2] pid=916738)�[0m Saved checkpoint to /lustre/fs1/portfolios/coreai/users/shuangy/src/nemo-rl-main/RL/results/Qwen8B-DAPO-Runs_qwen8b-fp8-dapo-deepgemm-cudagraph/tmp_step_50/policy/weights
�[36m(MegatronPolicyWorker[rank=2] pid=916738)�[0m [SHARON] GPU Memory after saving checkpoint: 24.89GB allocated, 24.99GB reserved
�[36m(VllmGenerationWorker pid=906501)�[0m INFO 09-04 05:46:10 [block_pool.py:321] Successfully reset prefix cache�[32m [repeated 14x across cluster]�[0m
�[36m(VllmGenerationWorker pid=906502)�[0m INFO 09-04 05:46:11 [gpu_worker.py:104] Sleep mode freed 47.40 GiB memory, 10.19 GiB memory is still in use.�[32m [repeated 7x across cluster]�[0m
�[36m(VllmGenerationWorker pid=906502)�[0m INFO 09-04 05:46:11 [executor_base.py:187] It took 1.054968 seconds to fall asleep.�[32m [repeated 7x across cluster]�[0m
�[36m(MegatronPolicyWorker[rank=1] pid=916742)�[0m [SHARON] GPU Memory before saving checkpoint: 24.90GB allocated, 25.06GB reserved�[32m [repeated 7x across cluster]�[0m
�[36m(MegatronPolicyWorker[rank=0] pid=916741)�[0m   successfully saved checkpoint from iteration       0 to /lustre/fs1/portfolios/coreai/users/shuangy/src/nemo-rl-main/RL/results/Qwen8B-DAPO-Runs_qwen8b-fp8-dapo-deepgemm-cudagraph/tmp_step_50/policy/weights [ t 1/1, p 1/2 ]
Removing checkpoint results/Qwen8B-DAPO-Runs_qwen8b-fp8-dapo-deepgemm-cudagraph/step_45 due to being outside top-3
Logged data to logs/grpo-qwen3-8b-1n8g-megatron-fp8-dapo/exp_035/train_data_step49.jsonl

📊 Training Results:
  • Loss: 0.0714
  • Avg Reward: 0.4794
  • Mean Generation Length: 3581.6812
  • Training FLOPS: 2617.26 TFLOPS (654.31 TFLOPS per rank)
  • Training Model Floating Point Utilization: 66.13%

⏱️  Timing:
  • Total step time: 2045.26s
  • generation: 1040.37s (50.9%)
  • policy_training: 577.73s (28.2%)
  • policy_and_reference_logprobs: 305.63s (14.9%)
  • checkpointing: 15.10s (0.7%)
  • data_processing: 5.59s (0.3%)
  • prepare_for_generation: 5.33s (0.3%)
  • training_prep: 1.68s (0.1%)
  • logprob_inference_prep: 0.61s (0.0%)
  • prepare_for_generation/transfer_and_update_weights: 0.52s (0.0%)
  • reward_calculation: 0.36s (0.0%)

========================= Step 51/300 =========================
▶ Preparing batch...
▶ Generating responses for batch of size 8192...
�[36m(VllmGenerationWorker pid=906504)�[0m INFO 09-04 05:47:12 [executor_base.py:203] It took 0.298370 seconds to wake up tags {'weights', 'kv_cache'}.
�[36m(MegatronPolicyWorker[rank=1] pid=916742)�[0m Saved checkpoint to /lustre/fs1/portfolios/coreai/users/shuangy/src/nemo-rl-main/RL/results/Qwen8B-DAPO-Runs_qwen8b-fp8-dapo-deepgemm-cudagraph/tmp_step_50/policy/weights�[32m [repeated 7x across cluster]�[0m
�[36m(MegatronPolicyWorker[rank=1] pid=916742)�[0m [SHARON] GPU Memory after saving checkpoint: 24.90GB allocated, 25.06GB reserved�[32m [repeated 7x across cluster]�[0m
�[36m(VllmGenerationWorker pid=906503)�[0m INFO 09-04 05:47:12 [executor_base.py:203] It took 0.348215 seconds to wake up tags {'kv_cache', 'weights'}.
�[36m(VllmGenerationWorker pid=906504)�[0m 
Adding requests:   0%|          | 0/1024 [00:00<?, ?it/s]
�[36m(MegatronPolicyWorker[rank=3] pid=916743)�[0m WARNING:megatron.core.dist_checkpointing.validation:Rank 3 common state dict differs from rank 0 common state dict. Keys only on rank 0: [], Keys only on 3: [], Mismatched keys: [(('optimizer', 'optimizer', 'param_groups', 1, 'step'), <class 'int'>, <class 'int'>)]
�[36m(VllmGenerationWorker pid=906504)�[0m 
Adding requests: 100%|██████████| 1024/1024 [00:00<00:00, 11827.08it/s]
�[36m(VllmGenerationWorker pid=906504)�[0m 
Processed prompts:   0%|          | 0/1024 [00:00<?, ?it/s, est. speed input: 0.00 toks/s, output: 0.00 toks/s]
Traceback (most recent call last):
  File "/lustre/fs1/portfolios/coreai/users/shuangy/src/nemo-rl-main/RL/./examples/run_grpo_math.py", line 276, in <module>
    main()
  File "/lustre/fs1/portfolios/coreai/users/shuangy/src/nemo-rl-main/RL/./examples/run_grpo_math.py", line 259, in main
    grpo_train(
  File "/lustre/fs1/portfolios/coreai/users/shuangy/src/nemo-rl-main/RL/nemo_rl/algorithms/grpo.py", line 600, in grpo_train
    repeated_batch, rollout_metrics = run_multi_turn_rollout(
                                      ^^^^^^^^^^^^^^^^^^^^^^^
  File "/lustre/fs1/portfolios/coreai/users/shuangy/src/nemo-rl-main/RL/nemo_rl/experience/rollouts.py", line 404, in run_multi_turn_rollout
    active_batch, generated_ids, gen_metrics = generate_responses(
                                               ^^^^^^^^^^^^^^^^^^^
  File "/lustre/fs1/portfolios/coreai/users/shuangy/src/nemo-rl-main/RL/nemo_rl/experience/rollouts.py", line 67, in generate_responses
    generation_outputs = policy_generation.generate(
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lustre/fs1/portfolios/coreai/users/shuangy/src/nemo-rl-main/RL/nemo_rl/models/generation/vllm/vllm_generation.py", line 387, in generate
    results = self.worker_group.get_all_worker_results(future_bundle)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lustre/fs1/portfolios/coreai/users/shuangy/src/nemo-rl-main/RL/nemo_rl/distributed/worker_groups.py", line 903, in get_all_worker_results
    return future_bundle.get_results(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lustre/fs1/portfolios/coreai/users/shuangy/src/nemo-rl-main/RL/nemo_rl/distributed/worker_groups.py", line 99, in get_results
    all_results = ray.get(object_refs)
                  ^^^^^^^^^^^^^^^^^^^^
  File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/auto_init_hook.py", line 21, in auto_init_wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/client_mode_hook.py", line 103, in wrapper
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/worker.py", line 2822, in get
    values, debugger_breakpoint = worker.get_objects(object_refs, timeout=timeout)
                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/worker.py", line 930, in get_objects
    raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(OutOfMemoryError): �[36mray::VllmGenerationWorker.generate()�[39m (pid=906502, ip=10.65.0.143, actor_id=a58c0e8e6fc5c3861fdf858f01000000, repr=VllmGenerationWorker)

As you could see in the log, I printed out the memory utilization before and after checkpoing saving:

�[36m(MegatronPolicyWorker[rank=2] pid=916738)�[0m [SHARON] GPU Memory before saving checkpoint: 24.89GB allocated, 24.99GB reserved
�[36m(MegatronPolicyWorker[rank=2] pid=916738)�[0m [SHARON] GPU Memory after saving checkpoint: 24.89GB allocated, 24.99GB reserved

I think the reason is for the megatron path, optimizer states is not freed up after saving the checkpoint.

I tried by adding a logic to move the optimizer states to cpu which resolves the problem for me:

    def save_checkpoint(
        self,
        weights_path: str,
        optimizer_path: Optional[str] = None,
        **kwargs,
    ):
        """Save a training checkpoint.

        Args:
            weights_path: The specific directory path where the checkpoint will be saved.
            optimizer_path: If not None, optimizer and scheduler states are saved if they exist.
        """
        allocated = torch.cuda.memory_allocated() / (1024**3)  # Convert to GB
        reserved = torch.cuda.memory_reserved() / (1024**3)  # Convert to GB
        print(f"[SHARON] GPU Memory before saving checkpoint: {allocated:.2f}GB allocated, {reserved:.2f}GB reserved")
        if not torch.distributed.is_initialized():
            raise RuntimeError(
                "Distributed process group is not initialized. Cannot save checkpoint."
            )

        if self.mcore_state is None or self.model is None:
            raise RuntimeError(
                "Megatron core state or model is not initialized. Cannot save checkpoint."
            )

        original_save_path = self.mcore_state.cfg.checkpoint_config.save
        # save_dir = os.path.dirname(weights_path)
        release_name = os.path.basename(weights_path)

        try:
            maybe_finalize_async_save(
                ckpt_cfg=self.mcore_state.cfg.checkpoint_config, blocking=False
            )
            self.mcore_state.cfg.checkpoint_config.save = weights_path

            optimizer_to_save = None
            scheduler_to_save = None

            if optimizer_path is not None:
                if self.optimizer is not None:
                    optimizer_to_save = self.optimizer
                if self.scheduler is not None:
                    scheduler_to_save = self.scheduler

            # Ensure model is in eval mode for consistent saving, unless actively training
            # This is a common practice, though NeMo's save might handle this.
            # For safety, if not in training loop, setting to eval.
            is_training = self.model.training
            if not is_training:
                self.model.eval()

            if self.should_disable_forward_pre_hook:
                self.disable_forward_pre_hook()
            save_checkpoint(
                state=self.mcore_state,
                model=[self.model],
                optimizer=optimizer_to_save,
                opt_param_scheduler=scheduler_to_save,
                num_floating_point_operations_so_far=self.mcore_state.train_state.floating_point_operations_so_far,
                checkpointing_context=self.checkpointing_context,
            )
            print(f"Saved checkpoint to {weights_path}")
            maybe_finalize_async_save(
                ckpt_cfg=self.mcore_state.cfg.checkpoint_config,
                blocking=True,
                terminate=True,
            )
            if self.should_disable_forward_pre_hook:
                self.enable_forward_pre_hook()

            if not is_training:  # Restore training state if it was changed
                self.model.train()
            
            # Move optimizaer states to cpu
            torch.randn(1).cuda()  # wake up torch allocator
            if hasattr(self, "optimizer") and self.optimizer is not None:
                # Iterate through the state dictionaries for each parameter group
                if isinstance(self.optimizer, ChainedOptimizer):
                    optimizer_state = self.optimizer.state
                else:
                    optimizer_state = self.optimizer._get_state()
                for _, state in optimizer_state.items():
                    # Iterate through the state items (e.g., momentum, variance) for a parameter
                    for k, v in state.items():
                        # Check if the item is a tensor and on the GPU
                        if torch.is_tensor(v) and v.is_cuda:
                            # Move the tensor to CPU and update the state dictionary
                            state[k] = v.to("cpu")

            gc.collect()
            torch.cuda.empty_cache()

            allocated = torch.cuda.memory_allocated() / (1024**3)  # Convert to GB
            reserved = torch.cuda.memory_reserved() / (1024**3)  # Convert to GB
            print(f"[SHARON] GPU Memory after saving checkpoint: {allocated:.2f}GB allocated, {reserved:.2f}GB reserved")


        except Exception as e:
            print(f"Failed to save checkpoint to {weights_path}: {e}")
            raise
        finally:
            self.mcore_state.cfg.checkpoint_config.save = original_save_path

[zhandaz]

I have tested with 4x8 H100s on Qwen3-30B-A3B with TP4PP8. Previously, someone reported OOM issues during checkpoint saving in this case, but this is no longer happening. I think this PR #1159 has resolved the issue. Please feel free to reopen this issue if it still exists.

cc: @wangshangsam @terrykong @yaoyu-33