Fix FusedAdam empty tensor handling

[wujingyue]

Summary

• Filter zero-sized optimizer tensor slots through a shared helper before Adam/SGD call multi_tensor_apply
• Add a native multi_tensor_apply check so zero-chunk tensors fail with a clear error if they reach the launcher
• Add regression tests for an empty parameter at the end of Adam and SGD parameter groups

Root cause

multi_tensor_apply drives launch bookkeeping from the per-tensor chunk loop. A zero-sized tensor has zero chunks, so if an empty tensor appears at the end of a group, pending work for earlier tensors can remain unlaunched and FusedAdam silently skips the update.

The generic MultiTensorApply wrapper still preserves caller-visible slot semantics. Adam and SGD opt into filtering because empty tensor slots are no-op optimizer updates.

Fixes #3207

Validation

• NVTE_FRAMEWORK=pytorch NVTE_CUDA_ARCHS=86 NVTE_WITH_NCCL_EP=0 NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD=1 MAX_JOBS=4 NVTE_BUILD_THREADS_PER_JOB=1 python setup.py build_ext --inplace
• python -m py_compile transformer_engine/pytorch/optimizers/fused_adam.py transformer_engine/pytorch/optimizers/fused_sgd.py transformer_engine/pytorch/optimizers/multi_tensor_apply.py
• PYTHONPATH=/tmp/TransformerEngine pytest -q tests/pytorch/test_fused_optimizer.py::TestFusedAdam::test_empty_param_at_end_of_group tests/pytorch/test_fused_optimizer.py::TestFusedSGD::test_empty_param_at_end_of_group
• All-empty sanity for te.optimizers.FusedAdam and te.optimizers.FusedSGD: both completed step() and torch.cuda.synchronize()
• Issue reproducer: te 0.0009999871253967285, torch 0.0009999871253967285
• git diff --check

[greptile-apps]

Greptile Summary

This PR fixes empty tensor handling in fused optimizer multi-tensor launches. The main changes are:

• Empty optimizer slots are filtered before Adam and SGD launch work.
• The native multi-tensor launcher now reports a clear error for unfiltered zero-sized tensors.
• Adam and SGD tests cover an empty parameter at the end of a group.

Confidence Score: 5/5

This looks safe to merge.

• No blocking issues found in the changed code.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/optimizers/multi_tensor_apply.py	Adds aligned empty-slot filtering before dispatching to the native multi-tensor op.
transformer_engine/pytorch/optimizers/fused_sgd.py	Routes SGD launch sets through the shared filtering path before native launch.
transformer_engine/common/multi_tensor/multi_tensor_apply.cuh	Adds a native guard for zero-chunk tensors that reach the shared launcher.
tests/pytorch/test_fused_optimizer.py	Adds Adam and SGD coverage for empty parameters at the end of a group.

_{Reviews (6): Last reviewed commit: "Move empty tensor filtering into MultiTe..." | Re-trigger Greptile}

[wujingyue]

Should we call this from MultiTensorApply unconditionally? I'm not sure whether filtering is safe for every optimizer.

[vthumbe1503]

Yes, agreed, I dont think there is any harm applying this filter directly in the MultiTensorApply class instead of putting it in each optimizer. Could we make that change @wujingyue ?

[vthumbe1503]

Non blocking comment, given that we have just have optimizers exposed in TE

[wujingyue]

Done

[vthumbe1503]

LGTM. Just one minor review comment

[vthumbe1503]

/te-ci pytorch

[vthumbe1503]

/te-ci

[wujingyue]

FYI, CI has an unrelated build error:

  nccl_ep.cc:23:10: fatal error: nccl_device.h: No such file or directory
     23 | #include <nccl_device.h>
        |          ^~~~~~~~~~~~~~~
  compilation terminated.