[PyTorch] Advance FusedAdam step counter for empty param groups

[adityasingh2400]

Fixes #1986

Root cause

FusedAdam.step() opens its param-group loop with an early skip for groups that hold no parameters:

for group in self.param_groups:
    if len(group["params"]) == 0:
        continue
    device = group["params"][0].device
    ...
    if "step" in group:
        group["step"] += ...
    else:
        group["step"] = ...

The step counter is updated after that skip, so an empty group never gets one. That is harmless for a group that is empty everywhere, but a group is often empty on only some data-parallel ranks. A no_weight_decay group holding just RMSNorm parameters is the usual case, and with pipeline or expert parallelism the ranks that own none of those parameters see an empty group while their peers do not.

step lives in param_groups, so state_dict() serializes it and it goes into the checkpoint. The ranks where the group was empty write step = null while the ranks where it was populated write the true iteration count, which is exactly what the counter table in the issue shows for a PP=2, EP=4, DP=8 run at iteration 2640. On resume, a rank that loads its optimizer shard from a rank where the group was empty picks up the stale value, and bias_correction computes 1 - beta1 ** step from a step that has nothing to do with how far training actually got.

Fix

Move the counter update above the empty-group skip so every group advances on every rank, then skip the kernel work for empty groups as before. This is the change suggested in the issue.

One detail the issue does not cover: with capturable=True the first update creates group["step"] as a device tensor and takes the device from group["params"][0], which an empty group does not have. The new code falls back to the device of self._dummy_overflow_buf, the optimizer's own scratch buffer, which is allocated on CUDA in __init__. Nothing else in the loop moved, so populated groups execute exactly the same sequence as before.

Verification

I do not have a GPU, so I could not run the TE test suite. Two things I did do.

The control flow itself is checked with a standalone CPU script that reproduces the loop head, before and after, on a real torch.optim.Optimizer so that param_groups and state_dict() behave as they do in TE. The kernel launch plays no part in the defect and is omitted:

import torch


class LoopHead(torch.optim.Optimizer):
    def __init__(self, params, fixed):
        super().__init__(params, {"lr": 1e-3})
        self.fixed = fixed

    def step(self):
        for group in self.param_groups:
            if self.fixed:
                # post-fix ordering
                if "step" in group:
                    group["step"] += 1
                else:
                    group["step"] = 1
                if len(group["params"]) == 0:
                    continue
            else:
                # ordering on main
                if len(group["params"]) == 0:
                    continue
                if "step" in group:
                    group["step"] += 1
                else:
                    group["step"] = 1


def run(fixed, num_steps=3):
    populated = torch.nn.Parameter(torch.zeros(4))
    optim = LoopHead([{"params": [populated]}, {"params": []}], fixed=fixed)
    for _ in range(num_steps):
        optim.step()
    return (
        [g.get("step") for g in optim.param_groups],
        [g.get("step") for g in optim.state_dict()["param_groups"]],
    )


for label, fixed in (("main", False), ("fixed", True)):
    print(label, run(fixed))

Output:

main  ([3, None], [3, None])
fixed ([3, 3], [3, 3])

The None in the checkpoint on main is the null step reported in the issue.

The regression test in this PR is the GPU version of the same property. test_empty_param_group_advances_step builds a FusedAdam over one populated group and one empty group, steps three times, and asserts that both groups report the same step in param_groups and in state_dict(). It is parametrized over capturable so the tensor-valued counter and the new device fallback are both exercised. On main the test fails at the first assertion on the empty group with a KeyError for step.

The changed files were formatted with the repository's pinned black 24.4.2 and the pre-commit arguments, and both are unchanged by it.

[greptile-apps]

Greptile Summary

The PR advances FusedAdam’s per-group step counter before skipping empty parameter groups, keeping serialized counters synchronized across distributed ranks. It also initializes capturable counters for empty groups on the optimizer scratch-buffer device and adds regression coverage for ordinary and capturable operation.

Confidence Score: 5/5

The PR appears safe to merge, with the empty-group counter invariant corrected and directly covered by the existing optimizer test suite.

The changed ordering advances every parameter group’s counter while preserving the existing kernel skip for empty groups, and the capturable fallback uses the optimizer’s CUDA scratch-buffer device without affecting populated-group execution.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/optimizers/fused_adam.py	Moves step bookkeeping ahead of the empty-group early return and supplies a CUDA-device fallback for capturable counters.
tests/pytorch/test_fused_optimizer.py	Adds a three-step regression test verifying empty-group counters in live optimizer state and serialized state for both capturable modes.

_{Reviews (1): Last reviewed commit: "[PyTorch] Advance FusedAdam step counter..." | Re-trigger Greptile}