NVIDIA / NVIDIA/Megatron-LM

[BUG] decoder_last_pipeline_num_layers interaction with PP=1

Open
#1,778 0 comments 1 reaction 1 assignee Claimed by @gautham-kollu View on GitHub
bug
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:

Image

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.