NVIDIA / NVIDIA/Megatron-LM

[Megatron Lite] CP+PP eval-only BSHD forward passes unsliced labels to cross entropy

Open
#5,617 2 comments 0 reactions 1 assignee Claimed by @FDecaYed View on GitHub
bug waiting-on-maintainers
Dominant language
Python
Stars
17.9k
Forks
4.5k
Avg merge
4d 6h
Merged PRs (30d)
271

Description

**Describe the bug**

In the experimental Megatron-Lite runtime, a Qwen3-MoE eval-only forward with context parallelism and pipeline parallelism enabled together (`CP=2`, `PP=2`) crashes in vocab-parallel cross entropy because the final pipeline stage has CP-local logits for 64 tokens but still receives full-sequence labels for 128 tokens.

The failure is optimizer-backend independent: it reproduces with both Megatron distributed optimizer (`dist_opt`) and PyTorch FSDP2.

Scope: this report is for the `use_thd=false` BSHD path. The THD protocol has separate packing/CP-splitting logic and is not claimed to be affected.

**Steps/Code to reproduce bug**

Environment used:

- 1 node / 8 GPUs
- Python 3.12
- Qwen3-30B-A3B-Instruct-2507 config/weights
- 2 model layers retained for the smoke test
- `TP=1, ETP=1, EP=1, CP=2, PP=2`
- sequence length 128, 2 microbatches

From a checkout containing `experimental/lite`:

```bash
export PYTHONPATH="$PWD/experimental/lite:/path/to/Megatron-LM"
export MEGATRON_LITE_DETERMINISTIC=1
export CUBLAS_WORKSPACE_CONFIG=:4096:8

torchrun --nnodes 1 --nproc_per_node 8 --master_port 29771 \
experimental/lite/examples/bench/correctness.py run \
--backend mlite \
--hf-path /path/to/Qwen3-30B-A3B-Instruct-2507 \
--model-name qwen3_moe \
--impl lite \
--tp 1 --etp 1 --ep 1 --pp 2 --cp 2 \
--steps 1 --num-microbatches 2 --seq-len 128 \
--seed 42 --truncate-layers 2 --same-data-across-dp \
--override-optimizer-json '{"use_precision_aware_optimizer":false}' \
--skip-weight-hash \
--impl-cfg-json '{"optimizer":"fsdp2"}' \
--output-json /tmp/fsdp2-cp2-pp2.json
```

The correctness runner performs eval-only forward before the training step, so the crash occurs immediately. Replacing the last argument with `{"optimizer":"dist_opt"}` produces the same failure.

Representative stack from FSDP2 (ranks 4-7, the last PP stage):

```text
correctness.py:338 in run_backend
eval_logits = _hash_tensor(_forward_logits(rt, handle, eval_batch))
runtime/backends/mlite/runtime.py:445 in forward_backward
outputs = forward_backward_pipelining(... forward_only=True)
primitive/parallel/pipeline.py:555 in _forward_only_pipeline_schedule
out = _run_pipeline_chunk_forward(...)
model/qwen3_moe/lite/protocol.py:137 in _forward_step_bshd
return model(input_ids=batch.input_ids.reshape(1, -1), labels=labels, ...)
model/qwen3_moe/lite/model.py:518 in forward
token_loss = vocab_parallel_cross_entropy(logits, labels_sb, self.ps.tp_group)
primitive/ops/cross_entropy.py:53 in forward
predicted_logits_1d = logits_2d[arange_1d, masked_target_1d]
IndexError: shape mismatch: indexing tensors could not be broadcast together with shapes [64], [128]
```

The apparent data flow is:

1. The bench supplies a raw model-agnostic `PackedBatch` with 128 input IDs and 128 labels.
2. With `PP>1`, `_infer_pipeline_tensor_shape()` sizes PP activations using the CP-local sequence length (`128 / 2 = 64`).
3. With `use_thd=false`, Qwen3-MoE selects `_forward_step_bshd()`, which reshapes and forwards the full `batch.labels` without applying the CP layout.
4. The final stage produces logits for 64 local tokens, while cross entropy receives 128 targets.

As controls, FSDP2 CP-only (`CP=2, PP=1`) and PP-only (`CP=1, PP=2`) smoke runs complete; the failure appears when CP and PP are enabled together in this eval-only path.

**Expected behavior**

Eval-only forward with `CP=2, PP=2` should apply the same CP token layout to labels/loss masks as to the final-stage activations, and should complete for both `dist_opt` and FSDP2.

At minimum, unsupported BSHD+CP combinations should fail early with a clear validation error rather than reaching cross entropy with inconsistent shapes.

**Additional context**

The relevant shared BSHD forward currently passes labels unchanged:

```python
def _forward_step_bshd(model: nn.Module, batch: PackedBatch) -> dict:
labels = batch.labels.reshape(1, -1) if batch.labels is not None else None
return model(input_ids=batch.input_ids.reshape(1, -1), labels=labels, packed_seq_params=None)
```

The THD forward instead calls `pack_thd_forward_kwargs()`, whose packing path CP-splits tokens, labels, masks, and position IDs. A likely fix is to give the BSHD path equivalent CP-aware label/loss-mask handling, or explicitly reject BSHD+CP until that layout is supported.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.