[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 counter before skipping empty parameter groups, keeping checkpointed group steps synchronized across distributed ranks.

• Uses the optimizer scratch buffer’s CUDA device when initializing a capturable counter for an empty group.
• Adds regression coverage for capturable and non-capturable counters and serialized optimizer state.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/optimizers/fused_adam.py	Moves group counter maintenance ahead of the empty-group skip and provides a CUDA-device fallback for capturable initialization.
tests/pytorch/test_fused_optimizer.py	Adds focused FusedAdam coverage verifying empty-group counters in live and serialized optimizer state.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A["FusedAdam.step()"] --> B["Iterate parameter group"]
  B --> C{"step exists?"}
  C -->|Yes| D["Increment step"]
  C -->|No| E["Initialize step"]
  D --> F{"group empty?"}
  E --> F
  F -->|Yes| G["Skip kernel work"]
  F -->|No| H["Prepare tensors and dispatch fused Adam kernel"]

Loading

_{Reviews (2): Last reviewed commit: "Merge branch 'main' into fix-fused-adam-..." | Re-trigger Greptile}

[ptrendx]

/te-ci pytorch

[adityasingh2400]

Flagging that the three red build jobs all died before any test ran, so none of them is exercising this change.

JAX      failure   failed step: Build
PyTorch  failure   no failed step, Build never completed, Sanity check never ran
All      failure   no failed step, same shape
Core     still in progress

JAX is the only one with a real failed step, and it is a CUDA link error rather than anything Python:

fatbinary fatal : Could not open input file 'flash_attn.compute_75.cubin'
error: subprocess-exited-with-error

PyTorch and All are stranger. Both report conclusion=failure with no failed step at all, their Build step still showing in_progress and Sanity check never reached. PyTorch started at 00:24:27 and completed at 02:45:31, so it sat in Build for 2h21m before the job was marked failed. That reads like a lost runner or a timeout during the docker build rather than a test result.

This PR is 21 lines of Python in fused_adam.py plus a test. It cannot make a .cubin fail to link or a container build hang for two hours, and in the PyTorch job the sanity check that would import it never executed.

Not asking for anything, just did not want the three reds to read as coming from the change while blossom-ci and te-ci are still pending. Happy to rebase onto current main if you think a fresh run would help.

[adityasingh2400]

Update on the run you triggered, since Core has now finished too.

JAX      failure    step 6 Build      died after 3m28s
PyTorch  failure    step 8 Build      still in_progress when killed, Sanity check never started
All      failure    step 8 Build      same shape
Core     cancelled  step 6 Build      cancelled at exactly 6h

All four died in Build, so still nothing here has run a test.

I also found the line just above the fatbinary error I quoted earlier, which points at the runner rather than at the code:

[  1%] Building CUDA object CMakeFiles/transformer_engine.dir/fused_attn/flash_attn.cu.o
fatbinary fatal   : Could not open input file 'flash_attn.compute_75.cubin'
sccache: Compiler killed by signal 1
gmake[2]: *** [CMakeFiles/transformer_engine.dir/build.make:76: ... flash_attn.cu.o] Error 254

The compiler was killed, so the cubin was never written. That is at 1 percent of the build, compiling transformer_engine/common/fused_attn/flash_attn.cu. This PR only touches transformer_engine/pytorch/optimizers/fused_adam.py and tests/pytorch/test_fused_optimizer.py, both Python, so it cannot reach that compile.

[adityasingh2400]

Correcting my last comment. I said nothing here had run a test. That was true of the four Build jobs, but te-ci has since reported and it did run the suite, so this change is now exercised.

test_fused_optimizer.py is green everywhere it ran:

A100   28 passed,  7 skipped
L40    30 passed,  5 skipped
B200   35 passed
H100   no pytest failures in the log

The red te-ci jobs are failing elsewhere. On A100 and L40 almost all of it is one environment assert, 590 of 680 on A100 and 594 on L40:

assert arch // 10 in [9, 10, 11, 12], "Unsupported compute capability. Supported: 9.x, 10.x, 11.x, 12.x"
E   AssertionError: Unsupported compute capability. Supported: 9.x, 10.x, 11.x, 12.x
/usr/local/lib/python3.12/dist-packages/flash_attn/cute/interface.py:1353

A100 is 8.0 and L40 is 8.9, so both land on 8 and fail that check. The installed flash_attn cute path does not cover those two cards. B200 is different, 41 failures across test_grouped_mlp.py, test_attention.py and test_kv_cache.py.

I could not account for the H100 and H100 debug statuses. Their uploaded logs show no pytest failures at all, so whatever turned them red is outside the part I can read, and I did not want to guess.

[ptrendx]

Rerunning CI with the updated branch.

[ptrendx]

/te-ci pytorch

[adityasingh2400]

Flagging that the rerun did not finish. PyTorch, All and Core were all cancelled at 00:26:46, so the pytorch suite never completed on 95a1af2.

JAX failed earlier, at 00:00:17, and it is a build error rather than a test one:

gmake[2]: *** [CMakeFiles/transformer_engine.dir/build.make:76:
CMakeFiles/transformer_engine.dir/fused_attn/flash_attn.cu.o] Error 254

This PR only touches transformer_engine/pytorch/optimizers/fused_adam.py and tests/pytorch/test_fused_optimizer.py, so that CUDA build is not coming from here. JAX C++ and JAX Python both passed on the same commit.

Happy to do anything on my side that helps, otherwise it just needs another run once the build is sorted.