[Dev] Fix expert-aware optimizer gradient statistics

[guapisolo]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

Ports the expert-aware optimizer gradient-statistics plumbing from #5916 to dev, on top of the duplicate-filter helper already added by #6165.

Issue tracking

Related to #5916, #6165, #6015, and #6099.

Problem

The standard optimizer creates separate dense and expert children, but both children only receive the dense tensor-parallel group. The gradient-norm and zero-count filters therefore deduplicate expert parameters with dense TP ownership.

For example, with TP4 / EP4 / ETP1, each rank owns distinct expert shards and every ETP group has local rank 0. The old filter instead keeps only dense TP rank 0, omitting valid expert gradients from ranks 1-3. Backward still produces those gradients and the optimizer still updates the parameters, but the reported global norm is too small, so clipping is weaker than intended.

Mixed-precision and distributed optimizers have a second part of the same issue: their parameter copies did not preserve allreduce=False, which is how the duplicate filter recognizes expert parameters.

Fix

• Attach both the dense TP and expert TP groups to standard optimizer wrappers.
• Preserve allreduce metadata when creating optimizer parameter copies or views.
• Pass both groups through gradient-norm and zero-count filtering, including the shared-group ChainedOptimizer path.

This is the optimizer-side subset of #5916. It intentionally does not port the separate gradient-synchronization, TE/native parameter-tagging, or parameter-norm logging changes from #5916, nor the LayerWise/Muon follow-up in #6099. As noted in #6015, this plumbing is also a prerequisite for a future dev counterpart of #6099.

Validation

• Added a distributed regression with TP=EP=world size and ETP1. One replicated dense scalar plus one distinct expert scalar per rank gives the exact oracle:
  • gradient norm: sqrt(1 + world_size);
  • zero count for all-zero gradients: 1 + world_size.
• Ran the new test with 8 ranks for both the standard BF16 optimizer and DistributedOptimizer: both passed on every rank.
• Ran tests/unit_tests/test_optimizer.py on one GPU: 28 passed, 42 skipped.
• Ran tools/autoformat.sh in check-only mode against dev: Black, isort, pylint, and Ruff passed (mypy reported the repository's existing ignored diagnostics).
• Additional Qwen3-30B-A3B integration validation of the same optimizer fix, using TP2 / CP2 / EP4 / ETP1:
  • two independent runs completed successfully;
  • the patched reported norm matched an independent FP64 logical reduction within 1.3e-7 relative error;
  • the old dense-TP counterfactual was 3.97% to 5.49% low across the four nonzero steps in the second run.

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests (the distributed unit test and external Qwen integration evidence cover this regression)
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation (no public API or user-facing behavior is added)
• I have run the autoformatter.sh on my PR

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[guapisolo]

Standalone reproduction on current dev

I reduced the Qwen integration failure to two scalars and the real optimizer path. This does not use Miles, a model checkpoint, or a helper-level mock: it runs DDP -> get_megatron_optimizer -> prepare_grads -> ChainedOptimizer.get_grad_norm/count_zeros.

The current official dev HEAD used below is 12eed0a5ac3f1ced6a3a4f76727db530bdbf7391. It already contains #6165 (de1ebd63c7935886150364f9fc1a4150c7bdb37c), so the duplicate-filter helper accepts an expert TP group. The reproduction shows that the optimizer still does not wire that group through, and mixed-precision optimizer parameters lose allreduce=False.

The topology is TP2 / EP2 / ETP1 / DP1. There is one replicated dense scalar and one distinct expert scalar per EP rank. With every gradient set to one, the exact norm is sqrt(1 + world_size) = sqrt(3). With every gradient set to zero, the exact zero count is 1 + world_size = 3. The affected dev path reports sqrt(2) and 2 instead.

repro_megatron_expert_grad_stats.py (SHA256: 729a7961c88be4010dbfcb65e927a26164b95e9dde3870755deb83075f4e2413)

#!/usr/bin/env python3
"""Reproduce expert-gradient statistic undercounting on Megatron-LM dev."""

import argparse
import hashlib
import inspect
import json
import math
import os
from pathlib import Path
import subprocess
import sys

import torch
import torch.distributed as dist
import torch.nn as nn


REPO_ROOT = Path.cwd().resolve()
sys.path.insert(0, str(REPO_ROOT))

from megatron.core import parallel_state  # noqa: E402
from megatron.core.distributed import (  # noqa: E402
    DistributedDataParallel,
    DistributedDataParallelConfig,
)
from megatron.core.optimizer import (
    OptimizerConfig,
    get_megatron_optimizer,
)  # noqa: E402
from megatron.core.transformer import TransformerConfig  # noqa: E402


class DenseAndExpertParameters(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.dense = nn.Parameter(torch.ones(1, dtype=torch.bfloat16, device="cuda"))
        self.expert = nn.Parameter(torch.ones(1, dtype=torch.bfloat16, device="cuda"))
        self.expert.allreduce = False


def scalar(value):
    return value.item() if isinstance(value, torch.Tensor) else value


def git_output(*args: str) -> str:
    return subprocess.check_output(
        ["git", *args], cwd=REPO_ROOT, text=True, stderr=subprocess.DEVNULL
    ).strip()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--expect", choices=("bug", "fixed"), required=True)
    parser.add_argument("--distributed-optimizer", action="store_true")
    args = parser.parse_args()

    world_size = int(os.environ["WORLD_SIZE"])
    rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])
    if world_size < 2:
        raise SystemExit("This reproduction requires at least two GPUs.")

    torch.cuda.set_device(local_rank)
    dist.init_process_group("nccl")
    parallel_state.initialize_model_parallel(
        tensor_model_parallel_size=world_size,
        expert_model_parallel_size=world_size,
        expert_tensor_parallel_size=1,
        create_gloo_process_groups=False,
    )

    report = None
    passed = False
    try:
        model = DistributedDataParallel(
            TransformerConfig(num_attention_heads=1, num_layers=1),
            DistributedDataParallelConfig(
                use_distributed_optimizer=args.distributed_optimizer
            ),
            DenseAndExpertParameters(),
        )
        optimizer = get_megatron_optimizer(
            OptimizerConfig(
                optimizer="adam",
                lr=0.0,
                bf16=True,
                clip_grad=0.0,
                log_num_zeros_in_grad=True,
                use_distributed_optimizer=args.distributed_optimizer,
            ),
            [model],
            use_gloo_process_groups=False,
        )

        dense_optimizer, expert_optimizer = optimizer.chained_optimizers
        for parameter in model.parameters():
            parameter.main_grad.fill_(1.0)
        if optimizer.prepare_grads():
            raise RuntimeError("prepare_grads reported non-finite gradients")
        actual_norm = float(scalar(optimizer.get_grad_norm()))

        for parameter in model.parameters():
            parameter.main_grad.zero_()
        if optimizer.prepare_grads():
            raise RuntimeError("prepare_grads reported non-finite gradients")
        actual_zero_count = int(scalar(optimizer.count_zeros()))

        expected_fixed_norm = math.sqrt(1 + world_size)
        expected_bug_norm = math.sqrt(2)
        expected_fixed_zero_count = 1 + world_size
        expected_bug_zero_count = 2
        tolerance = 1.0e-6

        fixed_match = (
            math.isclose(
                actual_norm, expected_fixed_norm, rel_tol=tolerance, abs_tol=tolerance
            )
            and actual_zero_count == expected_fixed_zero_count
        )
        bug_match = (
            math.isclose(
                actual_norm, expected_bug_norm, rel_tol=tolerance, abs_tol=tolerance
            )
            and actual_zero_count == expected_bug_zero_count
        )
        observed = (
            "FIX_VERIFIED"
            if fixed_match
            else "BUG_REPRODUCED" if bug_match else "UNEXPECTED"
        )
        wanted = "BUG_REPRODUCED" if args.expect == "bug" else "FIX_VERIFIED"
        passed = observed == wanted

        optimizer_source = Path(inspect.getsourcefile(get_megatron_optimizer)).resolve()
        if REPO_ROOT not in optimizer_source.parents:
            raise RuntimeError(
                f"Imported Megatron from {optimizer_source}, not requested checkout {REPO_ROOT}"
            )
        report = {
            "status": observed,
            "expectation": wanted,
            "passed": passed,
            "git_sha": git_output("rev-parse", "HEAD"),
            "git_status": git_output("status", "--porcelain"),
            "optimizer_source": str(optimizer_source),
            "optimizer_source_sha256": hashlib.sha256(
                optimizer_source.read_bytes()
            ).hexdigest(),
            "world_size": world_size,
            "topology": {"tp": world_size, "ep": world_size, "etp": 1, "dp": 1},
            "distributed_optimizer": args.distributed_optimizer,
            "actual_grad_norm": actual_norm,
            "expected_bug_grad_norm": expected_bug_norm,
            "expected_fixed_grad_norm": expected_fixed_norm,
            "actual_zero_count": actual_zero_count,
            "expected_bug_zero_count": expected_bug_zero_count,
            "expected_fixed_zero_count": expected_fixed_zero_count,
            "dense_tp_group_size": dense_optimizer.tp_group.size(),
            "expert_tp_group_size": (
                expert_optimizer.expert_tp_group.size()
                if hasattr(expert_optimizer, "expert_tp_group")
                else None
            ),
            "expert_optimizer_allreduce_metadata": [
                getattr(parameter, "allreduce", None)
                for parameter in expert_optimizer.get_parameters()
            ],
        }
        if rank == 0:
            print(json.dumps(report, sort_keys=True), flush=True)
    finally:
        dist.barrier()
        parallel_state.destroy_model_parallel()
        dist.destroy_process_group()

    if not passed:
        raise SystemExit(
            f"Reproduction expectation failed: {json.dumps(report, sort_keys=True)}"
        )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())

Run it from the root of a clean Megatron-LM checkout:

git fetch origin dev
git switch --detach 12eed0a5ac3f1ced6a3a4f76727db530bdbf7391

CUDA_VISIBLE_DEVICES=0,1 NCCL_NVLS_ENABLE=0 \
python -m torch.distributed.run --standalone --nproc_per_node=2 \
  /path/to/repro_megatron_expert_grad_stats.py --expect bug

# The same bug is present through DistributedOptimizer.
CUDA_VISIBLE_DEVICES=0,1 NCCL_NVLS_ENABLE=0 \
python -m torch.distributed.run --standalone --nproc_per_node=2 \
  /path/to/repro_megatron_expert_grad_stats.py \
  --expect bug --distributed-optimizer

For the PR head:

git fetch origin refs/pull/6470/head
git switch --detach FETCH_HEAD

CUDA_VISIBLE_DEVICES=0,1 NCCL_NVLS_ENABLE=0 \
python -m torch.distributed.run --standalone --nproc_per_node=2 \
  /path/to/repro_megatron_expert_grad_stats.py --expect fixed

CUDA_VISIBLE_DEVICES=0,1 NCCL_NVLS_ENABLE=0 \
python -m torch.distributed.run --standalone --nproc_per_node=2 \
  /path/to/repro_megatron_expert_grad_stats.py \
  --expect fixed --distributed-optimizer

Observed results:

checkout / path	actual norm	zero count	status
dev@12eed0a5 / BF16 optimizer, attempt 1	1.4142135623730951	2	BUG_REPRODUCED
dev@12eed0a5 / BF16 optimizer, attempt 2	1.4142135623730951	2	BUG_REPRODUCED
dev@12eed0a5 / DistributedOptimizer	1.4142135381698608	2	BUG_REPRODUCED
PR 3e72b966 / BF16 optimizer	1.7320508075688772	3	FIX_VERIFIED
PR 3e72b966 / DistributedOptimizer	1.7320507764816284	3	FIX_VERIFIED

On current dev, the report also shows expert_tp_group_size: null and expert_optimizer_allreduce_metadata: [null]. On the PR head these become 1 and [false], respectively. This isolates the remaining gap after #6165: the helper exists, but production optimizer group/metadata plumbing is missing.