[BUG] decoder_last_pipeline_num_layers interaction with PP=1
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Describe the bug**
When using `--pipeline-model-parallel-size 1` and `--decoder_last_pipeline_num_layers k`, Megatron does not detect when `k` does not match `num-layers` and does not raise an error. The model ends up effectively having `k` layers, without any warning of `num-layers` being overridden. In the logs, the number of parameters per rank is correct and corresponds to `k` layers, but the reported global number of parameters is computed based on `num-layers` and is wrong. Moreover, Megatron estimates FLOPs cost based on `num-layers` which leads to unrealistically high (and false) throughput. The load balancing seems to be summed over the `k` layers but then divided by `num-layers`, leading to strangely low values.
**To Reproduce**
E.g., Train a model with `--pipeline-model-parallel-size 1`, `--num-layers 20`, `--decoder_last_pipeline_num_layers 1`
**Expected behavior**
I believe `--decoder_last_pipeline_num_layers k` should either be ignored when PP=1, or there should be an assert verifying that `k` matches `num-layers` in this case.
**Stack trace/logs**
Not applicable
**Environment (please complete the following information):**
- Megatron-LM 650ab87d04105869f197f2ddc441e3b18ca93724
- PyTorch 25.06
- CUDA 12.9
- NCCL 2.27.3
**Proposed fix**
Assert that `decoder_last/first_pipeline_num_layers is None` or `pipeline-model-parallel-size > 1`
**Additional context**
An example of the absurd throughput you can get when Megatron trains a model with a single MoE block while thinking it has a lot more:
Contributor guide
Assessment
This issue has not been assessed yet.