[PyTorch] FusedAdam/multi_tensor_apply: OOB metadata writes when a zero-numel tensor lands in the last TensorListMetadata slot (CUDA IMA / misaligned address; silent corruption below threshold)

[yuhezhang-ai]

Describe the bug

multi_tensor_apply() (transformer_engine/common/multi_tensor/multi_tensor_apply.cuh)
gives every tensor a slot in the fixed-size kernel-arg struct
TensorListMetadata (36 slots for the depth-4 Adam kernel, 30 for the
depth-5 master-weights kernel). The tensors_full flush condition is only
evaluated inside the per-chunk loop. A zero-numel tensor occupies a
slot and increments loc_tensor_info, but contributes zero chunks, so the
chunk-loop body never runs for it. If the zero-numel tensor lands in the
last slot, loc_tensor_info reaches max_tensors without a flush, and
every subsequent tensor writes tl.sizes[loc_tensor_info] and
tl.addresses[d][loc_tensor_info] out of bounds of the struct arrays:
sizes[] writes overwrite block_to_tensor[], and the last addresses row
sprays raw pointer bytes over sizes[] and block_to_tensor[]. The launched
kernel then reads garbage tensor indices/sizes/addresses.

Consequences, depending on how many tensors follow the boundary-slot empty:

• ≥ ~16 trailing tensors (depth 5; ~20 for depth 4): CUDA fault —
  illegal memory access or misaligned address depending on the garbage
  bytes (we observed both signatures in production).
• Fewer trailing tensors: the overflow stays inside adjacent struct
  fields — no fault, silently wrong kernel inputs, i.e. corrupted
  optimizer updates with no error signal. This is arguably the worse case.

This input is not exotic: PyTorch FSDP2 (fully_shard) shards every
parameter along dim-0, so any parameter with dim-0 < shard-group size leaves
zero-numel local shards on tail ranks (e.g. [1, hidden] gates/scales,
class/position embeddings). Whether a training run crashes, silently
corrupts, or works is then decided by the numel-sequence-dependent slot
position of the empty shard — which is why the failure pattern looks
spurious (tail ranks only, first step only, only when small params are
trainable, "fixed" by unrelated param-count changes). We hit it
deterministically on 4-GPU and 128-GPU FSDP2 runs (VLM with trainable vision
tower; torch.optim AdamW on the identical population trains fine).

Steps/Code to reproduce the bug

Single GPU, no distributed setup, default kwargs:

import torch
from transformer_engine.pytorch.optimizers import FusedAdam

# 35 normal tensors, then a zero-numel tensor in the LAST metadata slot
# (36 slots for the depth-4 kernel), then enough tensors to push the OOB
# writes into block_to_tensor[].
params = [torch.nn.Parameter(torch.randn(8, device="cuda")) for _ in range(35)]
params.append(torch.nn.Parameter(torch.empty(0, device="cuda")))
params += [torch.nn.Parameter(torch.randn(8, device="cuda")) for _ in range(20)]
for p in params:
    p.grad = torch.randn_like(p)

opt = FusedAdam(params)  # plain kwargs; master-weights kwargs fault too (30-slot depth-5 kernel)
opt.step()               # CUDA illegal memory access / misaligned address
torch.cuda.synchronize()

Run with CUDA_LAUNCH_BLOCKING=1 to pin the fault to
multi_tensor_apply.cuh:92. For the depth-5 master-weights kernel use
[8]*29 + [0] + [8]*16 with
master_weights=True, master_weight_dtype=torch.float32, store_param_remainders=True, exp_avg_dtype=torch.bfloat16, exp_avg_sq_dtype=torch.float32.

Expected behavior

Zero-numel tensors are either skipped when building the tensor lists /
metadata (their update is an exact no-op — under FSDP2 every element lives on
other ranks; upstream apex skips empties in the analogous code), or rejected
with a clear Python-side error. The flush condition should in any case be
correct independent of chunk count (e.g. checked outside the chunk loop) so
that a slot-boundary empty cannot silently corrupt the metadata of subsequent
tensors.

Environment

• TE 2.15.0+42b84005 (source build), torch 2.12.0a0+nv26.04 (NGC 26.06),
  CUDA 13.2, H100. Fault verified deterministic single-GPU with the snippet
  above; also reproduced in real 4-GPU and 128-GPU FSDP2 trainings.
• The metadata-overflow structure is unchanged on current main (no
  numel()==0 guard; flush only inside the chunk loop).

Suggested fix

Skip numel() == 0 tensors when building TensorListMetadata (as upstream
apex does), and/or evaluate the tensors_full flush outside the per-chunk
loop; document that empty tensors are supported. Happy to test a fix.