Race condition in HF->mcore conversion when using megaton policy
- Dominant language
- Python
- Stars
- 2k
- Forks
- 562
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
If multiple runs are using the same shared directory for the megatron checkpoint dir, there is a potential race condition that may occur when:
1. One process starts the megatron conversion. During the conversion, it creates the checkpoint directory `os.path.join(pretrained_path, "iter_0000000")` https://github.com/NVIDIA-NeMo/RL/blob/f8fdb5c76915f196dd5220ac3ef572b2d52808e8/nemo_rl/models/policy/megatron_policy_worker.py#L425-L427
2. Before the first process finishes the megatron conversion, a second process sees the checkpoint directory and moves on, assuming the checkpoint exists https://github.com/NVIDIA-NeMo/RL/blob/f8fdb5c76915f196dd5220ac3ef572b2d52808e8/nemo_rl/models/policy/megatron_policy_worker.py#L434-L437
3. The second process will fail when actually trying to read the checkpoint because the first process hasn't finished writing it yet https://github.com/NVIDIA-NeMo/RL/blob/f8fdb5c76915f196dd5220ac3ef572b2d52808e8/nemo_rl/models/policy/megatron_policy_worker.py#L479-L482
Contributor guide
Research direction
Read nemo_rl/models/policy/megatron_policy_worker.py around lines 425-437 and 479-482, focusing on how the checkpoint directory is created, detected, and read. Reproduce or reason through two concurrent runs using the same shared directory. Done means a process cannot consume the checkpoint until the conversion has finished writing it.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100