NVIDIA / NVIDIA/Megatron-LM

Support variable-length packed sequences for seq_aux_loss

Open
#6,108 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.