🐛 Investigate large MoE num-zeros shifts after the 26.08 base-image update
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Describe the bug**
Investigate the large per-step gradient-zero shifts associated with the dev base-image update from `26.06-py3` to `26.08-py3` in [PR #6991: chore: Update base image to 26.08](https://github.com/NVIDIA/Megatron-LM/pull/6991).
The EP8/H100 and TP4/EP2/ETP2/PP2 one-node GB200 functional jobs failed `tests/functional_tests/python_test_utils/test_pretraining_regular_pipeline.py::test_regular_pipeline`. Their logs show that **both exact and approximate `num-zeros` comparisons failed, while approximate `lm loss` passed**. These are training-metric mismatches after training ran, not provisioning/checkout failures.
The PR now contains refreshed goldens. This issue tracks whether the gradient-zero changes are expected and reproducible; refreshing expected values alone does not establish their correctness. It does not claim that the original assertions still fail against the refreshed goldens.
@NVIDIA/mcore-oncall
**Failing run**
| Field | Value |
| --- | --- |
| PR | [#6991](https://github.com/NVIDIA/Megatron-LM/pull/6991) |
| Historical failing run | [33929057464](https://github.com/NVIDIA/Megatron-LM/actions/runs/33929057464), started 2026-09-04 UTC |
| Run branch / SHA | `pull-request/6991` / `3abfd97a00ff9574594dc61704b1dcb530956467` |
| EP8 H100 job | [`moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_optimizer - latest`](https://github.com/NVIDIA/Megatron-LM/actions/runs/33929057464/job/101220902508) |
| TP4/EP2/ETP2/PP2 one-node GB200 job | [`moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer_1node - latest`](https://github.com/NVIDIA/Megatron-LM/actions/runs/33929057464/job/101220899386) |
**Error**
The linked logs contain the following messages (timestamps, rank prefixes, and ANSI colors omitted):
```text
APPROXIMATE test for metric lm loss: PASSED
Exact comparison of num-zeros: FAILED
Approximate comparison of num-zeros: FAILED
AssertionError: The following metrics failed: lm loss, num-zeros, num-zeros
```
The logged checks use `ApproximateTest(atol=0, rtol=0.2)` for `num-zeros` and `ApproximateTest(atol=0, rtol=0.05)` for `lm loss`. Exact LM-loss comparison also failed; the approximate loss check passed. The assertions occur in multiple workload attempts in both linked jobs.
**Measured baseline shifts**
Comparison of the goldens at PR merge base `56ef6ed61ccb3edffa138ec581c212d479d9ec06` with PR head `a4d622074ccb990aaa4598923cec4c998774ea11`. This is a pinned PR-diff comparison, not a claim that all four rows failed in the same historical workflow.
Signed percentages use `100 × mean((old − new) / old)` over shared finite steps, excluding old values below `1e-12` in magnitude. All four rows have 100 contributing steps. These signed averages are not the CI pass/fail criterion.
| Test | Environment | `lm loss` (%) | `num-zeros` (%) | Mean zero count old → new |
| --- | --- | ---: | ---: | ---: |
| `gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_optimizer` | H100 | -0.000868% | -769.326429% | 102,872.59 → 64,326.41 |
| `gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer` | GB200 | -0.002910% | -138.533684% | 586,418.66 → 568,964.04 |
| `gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer` | H100 | -0.000189% | -95.279782% | 342,258.69 → 370,499.78 |
| `gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer_1node` | GB200 | -0.000914% | -35.784448% | 1,258,738.03 → 1,230,107.98 |
**Important interpretation:** `num-zeros` is the number of zero-valued gradient elements, not zero-valued weights. The reported percentage is an average of step-normalized changes, **not the change in the total or arithmetic-mean zero count**. Small old counts strongly affect this statistic. For example, EP8/H100 step 81 changes from **2,703 → 1,050,785** and step 70 from **6,590 → 1,051,568**, even though the overall arithmetic mean decreases. “769% more zeros overall” would therefore be incorrect.
Pinned golden files:
- [EP8/H100](https://github.com/NVIDIA/Megatron-LM/blob/a4d622074ccb990aaa4598923cec4c998774ea11/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_optimizer/golden_values_dev_dgx_h100.json)
- [TP4/EP2/ETP2/PP2 GB200](https://github.com/NVIDIA/Megatron-LM/blob/a4d622074ccb990aaa4598923cec4c998774ea11/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer/golden_values_dev_dgx_gb200.json)
- [TP4/EP2/ETP2/PP2 H100](https://github.com/NVIDIA/Megatron-LM/blob/a4d622074ccb990aaa4598923cec4c998774ea11/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer/golden_values_dev_dgx_h100.json)
- [TP4/EP2/ETP2/PP2 one-node GB200](https://github.com/NVIDIA/Megatron-LM/blob/a4d622074ccb990aaa4598923cec4c998774ea11/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp4_ep2_etp2_pp2_resume_torch_dist_dist_optimizer_1node/golden_values_dev_dgx_gb200.json)
**Steps/Code to reproduce bug**
1. Re-run the linked historical CI jobs using their recorded commit/image/configuration and pre-refresh expected values. Running against the current refreshed goldens is not an equivalent reproduction of the original assertion.
2. For a controlled investigation, hold source commit, data, seed, topology, and test configuration fixed; compare the same workload with the 26.06 and 26.08 environments and preserve per-step zero counts, loss, and routing/gradient statistics.
3. To reproduce the EP8/H100 summary statistic without launching training, with PR head `a4d622074ccb990aaa4598923cec4c998774ea11` already checked out and the comparison script's dependencies installed:
```bash
python tests/test_utils/python_scripts/compare_golden_values_kl.py \
--rev 56ef6ed61ccb3edffa138ec581c212d479d9ec06 \
--file tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_optimizer/golden_values_dev_dgx_h100.json
```
The script emits raw signed ratios; multiply `avg_rel_diff` by 100 for the percentages above. Repeat with the other three linked golden paths.
**Expected behavior / investigation needed**
- Explain whether the changed per-step gradient sparsity comes from expected numerical/routing changes, checkpoint-resume behavior, gradient-counting behavior, or an actual regression. **Root cause is not established** by the current evidence.
- Confirm reproducibility across repeats, including around checkpoint resume; a small mean LM-loss difference does not validate the gradient-zero changes by itself.
- Document the reason and evidence for accepting the new goldens, or correct the underlying issue and regenerate them if the shift is unintended.
**Additional context**
Investigated with the `mcore-create-issue` and `update-golden-values` skills. No workloads were rerun and no repository files or tolerances were changed while filing this issue. Assignment follows the latest author of the failing shared test file on `main` (`Phlip79`, commit `d2715d78a8b69261dc43cd0e9f7a2072788c32b3`); this is a triage assignment, not attribution of the regression.
Open issues were checked before filing. #5026/#5027 discuss older TE 2.16 NaNs/aux-kernel failures, and #5455 covers different weekly GPT FP8 tests; none describes these PR #6991 gradient-zero shifts.
Contributor guide
Assessment
This issue has not been assessed yet.