From 69337e75a75a88c134c20fc8b5638bc0e6467630 Mon Sep 17 00:00:00 2001 From: Jingyue Wu Date: Wed, 15 Jul 2026 01:09:48 +0000 Subject: [PATCH] Fix FusedAdam empty tensor handling Signed-off-by: Jingyue Wu --- tests/pytorch/test_fused_optimizer.py | 26 +++++++++++++++++++ .../multi_tensor/multi_tensor_apply.cuh | 3 +++ .../pytorch/optimizers/fused_adam.py | 6 ++++- .../pytorch/optimizers/fused_sgd.py | 6 +++-- .../pytorch/optimizers/multi_tensor_apply.py | 12 +++++++++ 5 files changed, 50 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/test_fused_optimizer.py b/tests/pytorch/test_fused_optimizer.py index a2863cba98..6832ef89dd 100644 --- a/tests/pytorch/test_fused_optimizer.py +++ b/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): diff --git a/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh b/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh index 6710a161b3..d6ac23d4af 100644 --- a/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh +++ b/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; diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index b48953e8e3..e664186ee2 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/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, diff --git a/transformer_engine/pytorch/optimizers/fused_sgd.py b/transformer_engine/pytorch/optimizers/fused_sgd.py index d7ab3fe9fe..9dfee73770 100644 --- a/transformer_engine/pytorch/optimizers/fused_sgd.py +++ b/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, diff --git a/transformer_engine/pytorch/optimizers/multi_tensor_apply.py b/transformer_engine/pytorch/optimizers/multi_tensor_apply.py index a5cbd27337..5a5b23e33a 100644 --- a/transformer_engine/pytorch/optimizers/multi_tensor_apply.py +++ b/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."""