Support variable-length packed sequences for seq_aux_loss
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
## Problem
Megatron Core's sequence-level MoE auxiliary loss needs the logical sample dimension, but compact THD packing flattens multiple samples into a single token stream.
The current `TransformerLayer._maybe_unflatten_for_moe()` path restores that dimension only when `PackedSeqParams.tokens_per_sample` is set. This works for equal-width flattened samples: `[mbs * S, 1, H]` can be reshaped to `[S, mbs, H]`.
Variable-length compact packs cannot provide one `tokens_per_sample` value. Their logical boundaries are represented by `cu_seqlens_q` / `cu_seqlens_q_padded`, and `tokens_per_sample` is necessarily `None`. The MoE router therefore sees `bsz=1`, causing `seq_aux_loss` to treat the entire flattened pack as one sequence rather than computing and averaging load balancing independently for each logical sample.
A `padding_mask` can correctly exclude physical alignment gaps from routing statistics, but it cannot by itself restore logical sequence ownership.
## Requested feature
Add boundary-aware `seq_aux_loss` support for variable-length packed sequences.
The implementation should:
- consume logical sequence ownership from packed-sequence boundaries (`cu_seqlens*`) or an equivalent explicit per-token sample index;
- compute sequence-level expert usage independently for each logical sample and average over the logical sample count;
- exclude physical alignment padding from expert counts and denominators;
- preserve the flattened THD tensor layout before and after the MoE layer;
- work with context parallelism and sequence parallelism, including CP-local token ordering/sharding;
- retain the existing `tokens_per_sample` fast path and its behavior for fixed-width packs.
A dense materialization to `[max_sequence_length, logical_batch_size, H]` may be undesirable for highly variable long-context packs, so a segmented/count-based implementation would be preferable if practical.
## Parallelism considerations
With context parallelism, each rank may hold non-contiguous portions of every logical sequence. Logical per-sequence expert counts and valid-token denominators must therefore be reduced over the appropriate TP/CP routing group without losing sample identity. The packed metadata may remain global while the token stream and padding mask are rank-local, so the API needs an unambiguous way to associate each local token with its original logical sequence.
## Acceptance criteria
- [ ] A variable-length pack (for example logical lengths `[3, 5]` with different physical padded boundaries) produces the same `seq_aux_loss` and gradients as an equivalent padded logical batch with a padding mask.
- [ ] Adding physical alignment gaps does not change the loss or expert counts.
- [ ] CP=2 produces parity with CP=1 for the same logical batch, including the current CP token ordering.
- [ ] Sequence-parallel execution remains correct.
- [ ] Forward output remains in the original flattened THD layout.
- [ ] Existing fixed-width `tokens_per_sample` behavior remains covered and unchanged.
- [ ] Tests cover forward loss, backward gradients, padding, and distributed CP behavior.
## Relevant code
- `megatron/core/packed_seq_params.py`
- `megatron/core/transformer/transformer_layer.py::_maybe_unflatten_for_moe`
- `megatron/core/transformer/moe/router.py::_apply_seq_aux_loss`
- Existing fixed-width support introduced by PR #5696
Contributor guide
Research direction
Start by reading megatron/core/packed_seq_params.py and TransformerLayer._maybe_unflatten_for_moe, then trace _apply_seq_aux_loss in megatron/core/transformer/moe/router.py and the fixed-width support from PR #5696. Run the existing fixed-width tests before adding coverage for variable-length padding, gradients, sequence parallelism, and CP=2 ordering. Done means parity with an equivalent padded logical batch while preserving flattened THD output and existing tokens_per_sample behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- distributed-systems, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100