NVIDIA-NeMo / NVIDIA-NeMo/RL

Race condition in HF->mcore conversion when using megaton policy

Open
#969 0 comments 0 reactions 0 assignees View on GitHub
bug
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.