[PyTorch] FusedAdam silently skips updates when a parameter group ends with an empty tensor

[wujingyue]

Describe the bug

transformer_engine.pytorch.optimizers.FusedAdam silently skips updates to
nonempty parameters when a parameter group ends with a zero-element parameter.
The optimizer step returns without an error, but the preceding nonempty parameter
remains unchanged. The equivalent torch.optim.Adam input updates correctly.

This matters for Megatron FSDP because it uses uneven parameter sharding. A valid
global parameter can therefore have a zero-element local shard on some
data-parallel ranks, making empty optimizer parameters common rather than an
invalid edge case. The resulting failure is rank-dependent and silently leaves
some shards stale.

Applications can work around this by detecting locally empty parameters and
excluding their gradients before every optimizer step. However, that adds
optimizer-specific filtering to every FSDP integration and complicates parameter
group, optimizer-state, and checkpoint bookkeeping. FusedAdam should accept a
valid empty local shard and treat its update as a no-op.

This appears related to #3202, which describes the underlying zero-numel handling
in multi_tensor_apply. This report captures a smaller two-parameter manifestation
that silently skips an otherwise valid update.

Steps/Code to reproduce the bug

Run the following on a single GPU; no distributed setup is required:

"""Reproduce TE FusedAdam skipping a parameter group ending in an empty tensor."""

import torch
import transformer_engine
from transformer_engine.pytorch.optimizers import FusedAdam


def run_optimizer(optimizer_type):
    """Run one Adam step over a nonempty parameter followed by an empty parameter."""
    parameters = [
        torch.nn.Parameter(torch.ones(4, device="cuda")),
        torch.nn.Parameter(torch.empty(0, device="cuda")),
    ]
    for parameter in parameters:
        parameter.grad = torch.ones_like(parameter)

    before = parameters[0].detach().clone()
    optimizer_type(parameters).step()
    return (parameters[0] - before).abs().max().item()


print(f"PyTorch: {torch.__version__}")
print(f"Transformer Engine: {transformer_engine.__version__}")
print(f"CUDA: {torch.version.cuda}")

te_change = run_optimizer(lambda params: FusedAdam(params, lr=1e-3))
torch_change = run_optimizer(lambda params: torch.optim.Adam(params, lr=1e-3))

print(f"TE FusedAdam max parameter change: {te_change}")
print(f"PyTorch Adam max parameter change: {torch_change}")

assert torch_change > 0, "The PyTorch Adam control did not update its nonempty parameter"
assert te_change > 0, (
    "TE FusedAdam silently skipped the nonempty parameter because the final parameter "
    "in the group was empty"
)

Observed output:

PyTorch: 2.12.0a0+0291f960b6.nv26.04.48445190
Transformer Engine: 2.16.0+4220403e
CUDA: 13.2
TE FusedAdam max parameter change: 0.0
PyTorch Adam max parameter change: 0.0009999871253967285
AssertionError: TE FusedAdam silently skipped the nonempty parameter because the final parameter in the group was empty

The final assertion is expected to fail and demonstrates the bug.

Expected behavior

FusedAdam should update every nonempty parameter in the group and ignore the empty
parameter as a no-op, matching torch.optim.Adam. It should not silently skip the
entire pending update because the final tensor is empty.

Environment overview

• Environment location: Docker, NVIDIA NGC PyTorch 26.04 development container
• Transformer Engine: preinstalled in the container

Environment details

• OS: Ubuntu 24.04.4 LTS
• Python: 3.12.3
• PyTorch: 2.12.0a0+0291f960b6.nv26.04.48445190
• Transformer Engine: 2.16.0+4220403e
• CUDA: 13.2
• cuDNN: 9.21

Device details

• GPU: NVIDIA RTX A6000
• Driver: 570.211.01