Fix FusedAdam empty tensor handling

tests/pytorch/test_fused_optimizer.py

@@ -166,6 +166,19 @@ def test_frozen_model(self):
 
             torch.testing.assert_close(ref_param, tst_param)
 
+    def test_empty_param_at_end_of_group(self):
+        tensors = [
+            torch.ones(4, dtype=torch.float, device="cuda"),
+            torch.empty(0, dtype=torch.float, device="cuda"),
+        ]
+        ref_param, tst_param, ref_optim, tst_optim = self.gen_param_optim(tensors, self.options)
+
+        self.gen_grad(ref_param, tst_param)
+        ref_optim.step()
+        tst_optim.step()
+
+        torch.testing.assert_close(ref_param, tst_param)
+
     def gen_precision_aware_test(
         self,
         use_fp8_params,
@@ -796,6 +809,19 @@ def test_float(self):
     def test_half(self):
         self.gen_single_type_test(param_type=torch.float16)
 
+    def test_empty_param_at_end_of_group(self):
+        tensors = [
+            torch.ones(4, dtype=torch.float, device="cuda"),
+            torch.empty(0, dtype=torch.float, device="cuda"),
+        ]
+        ref_param, tst_param, ref_optim, tst_optim = self.gen_param_optim(tensors, self.options)
+
+        self.gen_grad(ref_param, tst_param)
+        ref_optim.step()
+        tst_optim.step()
+
+        torch.testing.assert_close(ref_param, tst_param)
+
 
 class Model(torch.nn.Module):
     def __init__(self):

transformer_engine/common/multi_tensor/multi_tensor_apply.cuh

@@ -75,6 +75,9 @@ void multi_tensor_apply(int64_t block_size, int64_t chunk_size,
     loc_tensor_info++;
 
     auto chunks_this_tensor = (tensor_lists[0][t]->numel() + chunk_size - 1) / chunk_size;
+    NVTE_CHECK(chunks_this_tensor > 0,
+               "multi_tensor_apply expects tensors with at least one chunk; zero-sized tensors "
+               "must be filtered before launch because they skip the chunk loop");
 
     for (auto chunk = 0; chunk < chunks_this_tensor; chunk++) {
       tl.block_to_tensor[loc_block_info] = loc_tensor_info - 1;

transformer_engine/pytorch/optimizers/fused_adam.py

@@ -16,7 +16,7 @@
 from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor, Float8Quantizer
 from transformer_engine.pytorch.quantized_tensor import QuantizedTensor
 from ..constants import DType
-from .multi_tensor_apply import multi_tensor_applier
+from .multi_tensor_apply import filter_empty_tensor_lists, multi_tensor_applier
 
 
 def get_fp8_meta(fp8_tensor):
@@ -741,6 +741,10 @@ def apply_multi_tensor_adam(adam_func, tensor_lists, inv_scale=None, out_dtype=N
                 # pylint: disable=cell-var-from-loop
                 inv_scale_arg = () if inv_scale is None else (inv_scale,)
                 out_dtype_arg = () if out_dtype is None else (out_dtype,)
+                # Empty tensor slots are no-ops for Adam. If every slot was removed,
+                # skip the C++ wrapper because it reads tensor_lists[0][0].
+                if not filter_empty_tensor_lists(tensor_lists):
+                    return
                 multi_tensor_applier(
                     adam_func,
                     self._dummy_overflow_buf,

transformer_engine/pytorch/optimizers/fused_sgd.py

@@ -12,7 +12,7 @@
 from torch.optim.optimizer import Optimizer, required
 
 import transformer_engine_torch as tex
-from .multi_tensor_apply import multi_tensor_applier
+from .multi_tensor_apply import filter_empty_tensor_lists, multi_tensor_applier
 
 
 class FusedSGD(Optimizer):
@@ -295,7 +295,9 @@ def step(self, closure=None):
             for _, (launch_set, first_run) in enumerate(zip(launch_sets, first_runs)):
                 assert len(launch_set[0]) == len(launch_set[1])
                 assert len(launch_set[0]) == len(launch_set[2])
-                if len(launch_set[0]) > 0:
+                # Empty tensor slots are no-ops for SGD. If every slot was removed,
+                # skip the C++ wrapper because it reads tensor_lists[0][0].
+                if filter_empty_tensor_lists(launch_set):
                     multi_tensor_applier(
                         self.multi_tensor_sgd,
                         self._dummy_overflow_buf,

transformer_engine/pytorch/optimizers/multi_tensor_apply.py

@@ -6,6 +6,18 @@
 from torch.distributed._tensor import DTensor
 
 
+def filter_empty_tensor_lists(tensor_lists):
+    """Remove aligned zero-sized tensor slots and return whether any slots remain."""
+    if any(len(tensors) != len(tensor_lists[0]) for tensors in tensor_lists):
+        raise RuntimeError("Expected aligned multi-tensor lists.")
+
+    keep_slot = [tensor.numel() > 0 for tensor in tensor_lists[0]]
+    for i, tensors in enumerate(tensor_lists):
+        tensor_lists[i] = [tensor for tensor, keep in zip(tensors, keep_slot) if keep]
+
+    return bool(tensor_lists[0])
+
+
 class MultiTensorApply:  # pylint: disable=too-few-public-methods
     """Multi-tensor apply entry."""