NVIDIA-NeMo / NVIDIA-NeMo/RL

Async checkpoint saving in GRPO loop

Open
#1,471 1 comment 1 reaction 1 assignee Assigned to @parthmannan View on GitHub
enhancement Performance research t-checkpoint t-mcore
Dominant language
Python
Stars
2k
Forks
561
Avg merge
4d 5h
Merged PRs (30d)
145

Description

**Is your feature request related to a problem? Please describe.**
Currently async_save is [disabled](https://github.com/NVIDIA-NeMo/RL/blame/8762f575c0d11aeb8a073a64e49cf433eb77c94a/nemo_rl/models/policy/megatron_policy_worker.py#L654) in mcore path checkpoint, serialization takes a long time with training paused; should test async_save and resolve any issues coming up.

**Describe the solution you'd like**
A clear and concise description of what you want to happen.

**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.

**Additional context**
Add any other context or screenshots about the feature request here.

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.