several fixes for THD e2e - #5535
Conversation
Signed-off-by: xiaoyao0115 <1804647152@qq.com>
0b2ea65 to
9b4db6a
Compare
|
/ok to test 9b4db6a |
|
/claude strict-review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| """ | ||
|
|
||
| pad_packed_seq_by_appending_dummy_seq: bool = True | ||
| pad_packed_seq_by_appending_dummy_seq: bool = False |
There was a problem hiding this comment.
[IMPORTANT Compatibility] Changing the default from True to False is a silent behavior change for existing training scripts that rely on the default. Users upgrading without modifying their configs will get different cu_seqlens metadata (the padding tail now extends the last padded endpoint instead of creating a dummy sequence).
This is the correct long-term default (it aligns with the separation of valid vs. padded boundaries), but consider adding a deprecation note in the release/changelog so users relying on the old default know to set pad_packed_seq_by_appending_dummy_seq=True explicitly if needed.
| values = tracker["loss_sums"] | ||
| if tracker.get('reduce_group') is not None: | ||
| torch.distributed.all_reduce(values, group=tracker['reduce_group']) | ||
| if tracker.get('avg_group') is not None: | ||
| torch.distributed.all_reduce( | ||
| values, group=tracker['avg_group'], op=torch.distributed.ReduceOp.AVG | ||
| ) | ||
| tracker["values"] = values |
There was a problem hiding this comment.
[SUGGESTION Simplification] values = tracker["loss_sums"] is a reference, not a copy. The in-place all_reduce modifies tracker["loss_sums"], and then tracker["values"] = values makes both keys point to the same tensor. This means clean_loss_in_tracker() zeroing loss_sums also destroys values.
Currently safe because track_mtp_metrics calls _report (which reads values) before clean_loss_in_tracker. But the aliasing is non-obvious — a future caller reading tracker["values"] after cleanup would silently get zeros. Consider cloning to keep the tensors independent:
tracker["values"] = values.clone()
| tracker = MTPLossLoggingHelper.tracker | ||
| if "loss_sums" not in tracker: |
There was a problem hiding this comment.
[IMPORTANT Compatibility] This changes MTP logging semantics from global token-weighted average (sum(loss) / sum(tokens) across microbatches and ranks) to mean of per-microbatch normalized losses. The corresponding mtp_loss_scale in training.py changed from 1.0 to 1/num_microbatches.
For ongoing training runs, the logged MTP loss metric value will change at the upgrade boundary even without any real model change — the two formulas only agree when all microbatches have the same token count. This is documented as intentional ("preserves sequence-packing semantics"), but anyone comparing pre- and post-upgrade loss curves should be aware the metric definition changed.
| if has_dummy_padding_seq: | ||
| cu_seqlens_q = _append_dummy_seq(cu_seqlens_q, global_target_len) | ||
| cu_seqlens_kv = _append_dummy_seq(cu_seqlens_kv, global_target_len) | ||
| cu_seqlens_q_padded = _append_dummy_seq(cu_seqlens_q_padded, global_target_len) | ||
| cu_seqlens_kv_padded = _append_dummy_seq(cu_seqlens_kv_padded, global_target_len) | ||
| # None is TE's auto-detect mode. Keep real and physical boundaries | ||
| # separate unless the caller explicitly disables padding between sequences. | ||
| if ( | ||
| packed_seq_params.pad_between_seqs is not False | ||
| and cu_seqlens_q is not None | ||
| and cu_seqlens_q_padded is not None | ||
| ): | ||
| cu_seqlens_q = _append_dummy_seq(cu_seqlens_q, int(cu_seqlens_q[-1].item())) | ||
| cu_seqlens_q_padded = _append_dummy_seq(cu_seqlens_q_padded, global_target_len) | ||
| else: | ||
| cu_seqlens_q = _append_dummy_seq(cu_seqlens_q, global_target_len) | ||
| cu_seqlens_q_padded = _append_dummy_seq(cu_seqlens_q_padded, global_target_len) | ||
| if ( | ||
| packed_seq_params.pad_between_seqs is not False | ||
| and cu_seqlens_kv is not None | ||
| and cu_seqlens_kv_padded is not None | ||
| ): | ||
| cu_seqlens_kv = _append_dummy_seq(cu_seqlens_kv, int(cu_seqlens_kv[-1].item())) | ||
| cu_seqlens_kv_padded = _append_dummy_seq(cu_seqlens_kv_padded, global_target_len) | ||
| else: | ||
| cu_seqlens_kv = _append_dummy_seq(cu_seqlens_kv, global_target_len) | ||
| cu_seqlens_kv_padded = _append_dummy_seq(cu_seqlens_kv_padded, global_target_len) |
There was a problem hiding this comment.
[SUGGESTION Simplification] The Q and KV branches are exact mirror copies. Consider extracting a small helper to remove the duplication:
def _append_pair(cu_valid, cu_padded, pad_between_seqs, global_target_len):
if (
pad_between_seqs is not False
and cu_valid is not None
and cu_padded is not None
):
return (
_append_dummy_seq(cu_valid, int(cu_valid[-1].item())),
_append_dummy_seq(cu_padded, global_target_len),
)
return (
_append_dummy_seq(cu_valid, global_target_len),
_append_dummy_seq(cu_padded, global_target_len),
)Then call it twice for Q and KV. Not a correctness issue, just reduces 23 lines to ~5.
| valid_tokens = (~padding_mask).unsqueeze(-1) | ||
| probs = probs * valid_tokens | ||
| routing_map = routing_map & valid_tokens | ||
|
|
||
| # Apply token dropping after removing padding so padding tokens cannot | ||
| # consume expert capacity. |
There was a problem hiding this comment.
[IMPORTANT Correctness — verified safe] Good fix — masking padding tokens from probs and routing_map before token dropping is critical: without this, padding tokens consume expert capacity slots and can displace valid tokens.
The placement after routing but before apply_router_token_dropping is correct. The softmax and top-k still include padding tokens, but this has negligible impact since padding hidden states are zero/random and don't systematically bias expert selection. The aux loss path independently handles padding via compute_routing_scores_for_aux_loss(... padding_mask=padding_mask) below.
Code Review SummaryCRITICAL: 0 | IMPORTANT: 2 | SUGGESTION: 2 | Positive: 1 Key Findings[IMPORTANT Compatibility] Default change for [IMPORTANT Compatibility] MTP logging semantics change ( [SUGGESTION Simplification] The Q/KV dummy-sequence branches in [SUGGESTION] In Correctness VerificationTraced the four main changes end-to-end:
Overall AssessmentRisk: Low-Medium. The implementation is well-structured with comprehensive test updates. The two compatibility changes (default flip, metric semantics) are intentional but should be documented in release notes. No correctness or performance issues found. |
2 participants
What does this PR do ?
This PR fixes padding handling for variable-length THD and SBHD training:
cu_seqlensseparate from physicalcu_seqlens_padded. Trailing alignment padding only extends the padded endpoint, whilepad_between_seqsis left for TE to infer.sum(loss) / sum(tokens)) to accumulated normalized per-microbatch losses.Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.