NVIDIA / NVIDIA/Megatron-LM

[BUG] Last rank exits with error when wandb accessing tracker file before rank 0 completes writing

Open
#1,653 0 comments 0 reactions 0 assignees View on GitHub
bug community-request
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Describe the bug**
When saving checkpoints with wandb enabled, the last rank's wandb artifact requires a tracker file that is written by the first rank (rank 0). If the shared file system is slow, the last rank may attempt to access the tracker file before rank 0 has finished writing it, causing the last rank to exit with an error.

**To Reproduce**
Steps to reproduce the behavior. The easier it is to reproduce the faster it will get maintainer attention.
1. Enable wandb and checkpoint save
2. Use a slow shared file system or Artificially add a sleep() delay in rank 0's iter_finalize_fn function
3. Run training with multiple ranks
4. Observe that the last rank fails when trying to access the tracker file

**Expected behavior**
The training should complete without errors. The last rank should wait for rank 0 to complete writing the tracker file before attempting to access it, or there should be proper synchronization to prevent race conditions.

**Stack trace/logs**
[rank15]: ValueError: Path is not a file: '/xxx/latest_checkpointed_iteration.txt'
[rank15]:     raise ValueError(f"Path is not a file: {local_path!r}")
[rank15]:   File "/usr/local/lib/python3.12/dist-packages/wandb/sdk/artifacts/artifact.py", line 1422, in add_file
[rank15]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank15]:     return method(self, *args, **kwargs)
[rank15]:   File "/usr/local/lib/python3.12/dist-packages/wandb/sdk/artifacts/_validators.py", line 255, in wrapper
[rank15]:     artifact.add_file(tracker_filename)
[rank15]:   File "Megatron-LM/megatron/training/wandb_utils.py", line 36, in on_save_checkpoint_success
[rank15]:     wandb_utils.on_save_checkpoint_success(checkpoint_name, get_checkpoint_tracker_filename(save_dir), save_dir, iteration)
[rank15]:   File "Megatron-LM/megatron/training/checkpointing.py", line 571, in wandb_finalize_fn
[rank15]:     wandb_finalize_fn()
[rank15]:   File "Megatron-LM/megatron/training/checkpointing.py", line 576, in save_checkpoint
[rank15]:     save_checkpoint(iteration, model, optimizer, opt_param_scheduler,

**Environment (please complete the following information):**
- Megatron-LM 878d65fe006b07ccef715ccc5ab521afb94e78a8
- PyTorch 2.7.0a0+79aa17489c.nv25.4
- CUDA 12.9
- NCCL 2.26.3

**Proposed fix**
Fixed in #1654 , thanks for reviewing it.

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.