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:

• Adds a shared helper to remove zero-sized tensor slots before optimizer launches.
• Applies the helper in FusedAdam and FusedSGD.
• Adds a native launcher check with a clearer error for unfiltered zero-sized tensors.
• Adds Adam and SGD tests for empty parameters 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 a shared helper for filtering aligned empty tensor slots before multi-tensor launches.
transformer_engine/pytorch/optimizers/fused_adam.py	Filters Adam tensor lists before dispatch and skips launches when all slots are empty.
transformer_engine/pytorch/optimizers/fused_sgd.py	Filters SGD launch sets before dispatch and skips launches when all slots are empty.
transformer_engine/common/multi_tensor/multi_tensor_apply.cuh	Adds a native check that rejects zero-chunk tensors with a clear error.
tests/pytorch/test_fused_optimizer.py	Adds Adam and SGD tests for an empty parameter at the end of a parameter group.

_{Reviews (3): Last reviewed commit: "Fix FusedAdam empty tensor handling" | Re-trigger Greptile}

[wujingyue]

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