[bugfix] fix megatron pp4 max_epochs #5432
Conversation
There was a problem hiding this comment.
Summary of Changes
Hello @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 primarily addresses a bug in Megatron's training process related to the accurate termination of training based on max_epochs, particularly in distributed environments. It refines the batch iteration and consumption logic to ensure precise epoch counting and completion. Additionally, the PR includes improvements to distributed batch handling and broadcasting, making the training process more robust. A conditional update to the peft library dependency is also introduced, along with a minor documentation update for a supported model.
Highlights
- Megatron Epoch Iteration Correction: The core of this bugfix addresses an issue where the training loop's epoch termination, specifically when
max_epochsis defined, might not behave as expected in Megatron's distributed training environment. The changes ensure that batches are correctly counted and consumed across epochs, leading to precise training duration. - Distributed Batch Handling Enhancements: Improvements were made to how batch data is handled and broadcasted across different stages in Megatron's pipeline parallelism. This includes explicitly setting non-applicable batch fields to
Noneand refining the logic for processingpadding_freebatches, enhancing the robustness and stability of distributed training. - PEFT Dependency Update for MoE: The
peftlibrary dependency has been updated. For configurations involving Mixture of Experts (MoE) models, a newer version (peft>=0.15) is now required, ensuring compatibility with the latest features or bug fixes inpeftthat are relevant to MoE setups. - Documentation Update for Qwen Model: The documentation for supported models has been updated to reflect that the
Qwen3-4B-Thinking-2507model is now officially supported.
Using Gemini Code Assist
The 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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
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 .gemini/ folder in the base of the repository. Detailed instructions can be found here.
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
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces a bugfix for max_epochs when using Megatron with pipeline parallelism. The core logic is updated in new_cyclic_iter to correctly handle the remaining batches in the last epoch. Additionally, it refactors and improves the handling of data batches for intermediate pipeline stages, making the code more robust and removing a previous hack. My review includes suggestions to improve code clarity and reduce duplication in the new logic.
swift/megatron/trainers/base.py
Outdated
| @@ -74,14 +74,15 @@ def initialize_megatron(*_args, **kwargs): | |||
| def new_cyclic_iter(iterable): | |||
| args = get_args() | |||
| i = 0 | |||
| n_batch = 0 | |||
| while True: | |||
| is_training = getattr(args, 'is_training', False) | |||
| if is_training: | |||
| logger.info(f'The training of Epoch {i} starts...') | |||
| if is_training and args.max_epochs and i >= args.max_epochs - 1: | |||
| it = iter(iterable) | |||
| num_batches = args.global_batch_size // (args.micro_batch_size * args.data_parallel_size) | |||
There was a problem hiding this comment.
For clarity, it would be better to rename num_batches to num_microbatches or a more descriptive name like num_microbatches_per_global_batch, as it's calculated using a function that returns the number of micro-batches per global batch, not the number of batches in an epoch. This would make the logic in this function easier to follow, especially in lines 85 and 88 where it's also used.
| num_batches = args.global_batch_size // (args.micro_batch_size * args.data_parallel_size) | |
| num_microbatches = args.global_batch_size // (args.micro_batch_size * args.data_parallel_size) |
swift/megatron/trainers/utils.py
Outdated
| else: | ||
| batch['input_ids'] = None | ||
| batch['labels'] = None | ||
| batch['attention_mask'] = None | ||
| batch['position_ids'] = None | ||
| batch['loss_scale'] = None |
There was a problem hiding this comment.
This block of code, which nullifies parts of the batch for intermediate pipeline stages, is very similar to lines 127-132. To reduce duplication and improve maintainability, consider using a loop over a list of keys.
For example:
else:
for key in ('input_ids', 'labels', 'attention_mask', 'position_ids', 'loss_scale'):
batch[key] = NoneSimilarly, the block at lines 127-132 could be made more concise with tuple unpacking:
else:
input_ids, labels, attention_mask, position_ids, loss_scale = (None,) * 5
No description provided.