[megatron] megatron remove num_samples#9522
Conversation
![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 simplifies batch preparation and data collation by removing the num_samples parameter and instead deriving it dynamically from the length of seq_lens when packed_seq_params is present. While this cleanup streamlines the code, the review highlights a critical regression in get_last_tokens where packed_seq_params.num_samples is still accessed. Additionally, the feedback suggests adding safety checks when popping seq_lens from the batch and when accessing it within packed_seq_params to prevent potential KeyError or AttributeError exceptions.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Jintao-Huang
force-pushed
the
megatron_remove_num_samples
branch
from
823caf5 to
2cd141b
Compare
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request refactors the tracking and calculation of num_samples across various Megatron trainers, primarily deriving it from packed_seq_params.seq_lens.shape[0] during padding-free training rather than passing it explicitly or storing it in batch dictionaries. It also introduces a new test file tests/megatron/test_opsd.py. A critical issue was identified in swift/megatron/trainers/rlhf_mixin.py where accessing packed_seq_params.seq_lens directly will cause an AttributeError if packed_seq_params is None (such as during standard training without packing). A code suggestion has been provided to safely fall back to labels.shape[0] in this scenario.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| per_token_logps = -output_tensor | ||
| loss_mask = labels != -100 | ||
| per_token_logps = per_token_logps * loss_mask | ||
| num_samples = packed_seq_params.seq_lens.shape[0] |
There was a problem hiding this comment.
When padding_free is False, packed_seq_params is None. Accessing packed_seq_params.seq_lens will raise an AttributeError: 'NoneType' object has no attribute 'seq_lens', causing an unconditional crash during standard training (e.g., standard DPO/KTO without packing). We should safely fall back to labels.shape[0] when packed_seq_params is None.
| num_samples = packed_seq_params.seq_lens.shape[0] | |
| num_samples = packed_seq_params.seq_lens.shape[0] if packed_seq_params is not None else labels.shape[0] |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request refactors the handling of num_samples across various Megatron trainers and utility functions. Instead of passing num_samples as an explicit parameter or retrieving it from batch dictionaries, it is now dynamically calculated using packed_seq_params.seq_lens.shape[0] or len(seq_lens). This simplifies function signatures and improves consistency. Additionally, test configurations are updated, and a new test file test_opsd.py is added. There are no review comments provided, so I have no feedback to address.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
2 participants
No description provided.