[megatron] fix save barrier #6653
Conversation
Summary of ChangesHello @Jintao-Huang, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request improves the reliability of model weight saving in a distributed training context. By introducing a distributed barrier after the training loop, it ensures that all participating processes synchronize their state before the saving operation completes. This prevents potential inconsistencies or incomplete saves that could arise in a multi-node training setup, thereby enhancing the overall robustness and data integrity of the training process. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds a dist.barrier() to synchronize processes after training in a distributed Megatron setup. This is a good practice to ensure all ranks have completed saving their model weights before the program proceeds or exits. However, I've identified a potential race condition with the current placement of the barrier that could lead to incomplete execution for certain ranks. My review includes a comment explaining this issue and suggesting a more robust placement for your consideration.
|
|
||
| try: | ||
| self.trainer.train(train_dataset, val_dataset, data_collator) | ||
| dist.barrier() # Ensure all weights are saved completely |
There was a problem hiding this comment.
While adding a barrier here correctly synchronizes processes after training and ensures weights are saved, its placement inside the try block could lead to a race condition. After passing the barrier, processes that are not the last_rank will execute the empty finally block and may exit, while the last_rank process is still running the visualization logic. If a job scheduler sees that most processes have terminated, it might kill the entire job, interrupting the visualization task on the last rank.
A more robust approach is to place the dist.barrier() after the finally block. This ensures all processes, including the last rank, complete their final tasks before synchronization and exit. As the finally block is not part of this diff, I cannot provide a direct code suggestion, but I recommend moving the barrier to improve the robustness of the shutdown sequence.
2 participants
No description provided.