NVIDIA-NeMo / NVIDIA-NeMo/RL

GRPO fails to save checkpoint during second run

Open
#1,637 2 comments 0 reactions 1 assignee Claimed by @yfw View on GitHub
bug
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

**Describe the bug**

When running deepseek grpo, it is able to save checkpoint for training continuation. After that, RL can successfully start training from checkpoint provided, but fails immediately when it has to save another checkpoint

**Steps/Code to reproduce bug**

Run the GRPO DSv3 until first checkpoint (step 10), stop training after checkpoint is saved. Restart training from provided checkpoint, wait for next validation step (step 20) to see it fails

**Expected behavior**

You should be able to save and restore form checkpoint several times

**Additional context**
```
Traceback (most recent call last):
File "/workspace/code/run_grpo_math.py", line 235, in
main()
File "/workspace/code/run_grpo_math.py", line 218, in main
grpo_train(
File "/opt/nemo-rl/nemo_rl/algorithms/grpo.py", line 1428, in grpo_train
policy.save_checkpoint(
File "/opt/nemo-rl/nemo_rl/models/policy/lm_policy.py", line 747, in save_checkpoint
ray.get(futures)
File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/auto_init_hook.py", line 22, 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 104, in wrapper
return func(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo_rl_venv/lib/python3.12/site-packages/ray/_private/worker.py", line 2882, 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 968, in get_objects
raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(AssertionError): ray::MegatronPolicyWorker.save_checkpoint() (pid=26034, ip=10.65.3.69, actor_id=d6c14005d9a00fbe52840b0d01000000, repr=MegatronPolicyWorker[rank=32])
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo-rl/nemo_rl/models/policy/megatron_policy_worker.py", line 2336, in save_checkpoint
save_checkpoint(
File "/opt/nemo-rl/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/src/megatron/bridge/training/checkpointing.py", line 552, in save_checkpoint
state_dict = generate_state_dict(
^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo-rl/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge/src/megatron/bridge/training/checkpointing.py", line 1068, in generate_state_dict
state_dict["optimizer"] = optimizer.sharded_state_dict(state_dict, **(optim_sd_kwargs or {}))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/optimizer/optimizer.py", line 1201, in sharded_state_dict
self._synchronize_steps()
File "/opt/nemo-rl/3rdparty/Megatron-LM-workspace/Megatron-LM/megatron/core/optimizer/optimizer.py", line 1401, in _synchronize_steps
assert len(steps) <= 1, f"steps: {steps}"
^^^^^^^^^^^^^^^
AssertionError: steps: [10, 20]
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.