diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py index 6ab5f901f73..8366e6ca87a 100644 --- a/tests/pytorch/distributed/run_ep.py +++ b/tests/pytorch/distributed/run_ep.py @@ -33,6 +33,13 @@ _ep_combine_raw, _ep_dispatch_raw, ) +from transformer_engine.pytorch.ops.fused.moe_ep import ( + FusedMoeEp, + _cudnn_megamoe_supported, + _get_megamoe_combine_format, + finalize_moe_ep_resources, + is_moe_fusion_supported, +) from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1" @@ -1216,6 +1223,19 @@ def test_runtime_buffer_config_mismatch(self): with self.assertRaisesRegex(ValueError, "combine_bwd_quant_recipe"): te_ops.MoeCombine(config, mxfp8_buffer)(expert_out, topk_idx) + def test_megamoe_combine_format_env(self): + old_value = os.environ.get("NVTE_MEGAMOE_MXFP8_COMBINE") + try: + os.environ["NVTE_MEGAMOE_MXFP8_COMBINE"] = "0" + self.assertEqual(_get_megamoe_combine_format(), "bf16") + os.environ["NVTE_MEGAMOE_MXFP8_COMBINE"] = "1" + self.assertEqual(_get_megamoe_combine_format(), "mxfp8") + finally: + if old_value is None: + os.environ.pop("NVTE_MEGAMOE_MXFP8_COMBINE", None) + else: + os.environ["NVTE_MEGAMOE_MXFP8_COMBINE"] = old_value + def _run_dispatch_combine_identity(self, *, mxfp8): """Route, apply top-k weights, and combine back to local token order.""" recipe = MXFP8BlockScaling() if mxfp8 else None @@ -1330,16 +1350,42 @@ def test_megamoe_mxfp8_numerics(self): @_mxfp8_align_test def test_megamoe_mxfp8_cuda_graph_matches_eager(self): + """Fused Sequential forward/backward graph replay matches eager execution.""" + if torch.cuda.get_device_capability() != (10, 7): + self.skipTest("FusedMoeEp CUDA graph test requires SM107") + if not _cudnn_megamoe_supported(): + self.skipTest("installed cuDNN frontend does not provide stateless training") + self._run_megamoe_mxfp8_cuda_graph_matches_eager( + glu_interleave_size=32, + expect_fused=True, + ) + + @_mxfp8_align_test + def test_unfused_megamoe_mxfp8_cuda_graph_matches_eager(self): """Unfused Sequential forward/backward graph replay matches eager execution.""" - self._run_megamoe_mxfp8_cuda_graph_matches_eager() + self._run_megamoe_mxfp8_cuda_graph_matches_eager( + glu_interleave_size=None, + expect_fused=False, + ) - def _run_megamoe_mxfp8_cuda_graph_matches_eager(self): + def _run_megamoe_mxfp8_cuda_graph_matches_eager( + self, + *, + glu_interleave_size, + expect_fused, + ): recipe = MXFP8BlockScaling() graph_model, graph_fc1, graph_fc2, _ = self._make_megamoe_model( recipe=recipe, + glu_interleave_size=glu_interleave_size, ) + fusion_supported = is_moe_fusion_supported(tuple(graph_model), recipe) + if expect_fused and not fusion_supported: + self.skipTest("current configuration does not support FusedMoeEp") + self.assertEqual(fusion_supported, expect_fused) eager_model, eager_fc1, eager_fc2, _ = self._make_megamoe_model( recipe=recipe, + glu_interleave_size=glu_interleave_size, ) eager_model.load_state_dict(graph_model.state_dict()) @@ -1368,6 +1414,17 @@ def reset_graphed_model(): self.addCleanup(reset_graphed_model) + forward_ops = graph_model._module_groups[0]._forward_ops + backward_ops = graph_model._module_groups[0]._backward_ops + fused_ops = [op for op, _ in forward_ops if isinstance(op, FusedMoeEp)] + if expect_fused: + self.assertEqual(len(forward_ops), 1) + self.assertEqual(len(backward_ops), 1) + self.assertEqual(len(fused_ops), 1) + self.assertIs(backward_ops[0][0], fused_ops[0]) + else: + self.assertFalse(fused_ops) + # Replace the capture-time contents while retaining captured addresses. with torch.no_grad(): static_tokens.copy_(torch.randn_like(static_tokens)) @@ -1408,6 +1465,14 @@ def reset_graphed_model(): eager_topk_weights, static_topk_idx, ) + if expect_fused: + eager_fused_ops = [ + op + for op, _ in eager_model._module_groups[0]._forward_ops + if isinstance(op, FusedMoeEp) + ] + self.assertEqual(len(eager_fused_ops), 1) + self.assertIs(fused_ops[0]._resource, eager_fused_ops[0]._resource) tolerances = {"rtol": 0.125, "atol": 0.25} torch.testing.assert_close(graph_out_snapshot, eager_out, **tolerances) @@ -1465,12 +1530,12 @@ def test_megamoe_delayed_wgrad_bf16(self): delay_wgrad_compute=True, ) - @_eager_test_include @_mxfp8_align_test def test_megamoe_delayed_wgrad(self): self._run_megamoe_vs_reference( quantization="mxfp8", delay_wgrad_compute=True, + require_fused=True, ) @_eager_test_include @@ -1481,13 +1546,13 @@ def test_megamoe_delayed_main_grad_bf16(self): delay_wgrad_compute=True, ) - @_eager_test_include @_mxfp8_align_test def test_megamoe_delayed_main_grad(self): self._run_megamoe_vs_reference( quantization="mxfp8", accumulate_into_main_grad=True, delay_wgrad_compute=True, + require_fused=True, ) def _run_megamoe_vs_reference( @@ -1497,12 +1562,15 @@ def _run_megamoe_vs_reference( accumulate_into_main_grad=False, overwrite_main_grad=False, delay_wgrad_compute=False, + require_fused=False, ): """Compare the five-op MoE sequence with the PyTorch EP reference. The fuser selects MegaMoE when its runtime gates pass. Otherwise this exercises the same sequence as separate NCCL EP and grouped-MLP ops. """ + if require_fused and not _cudnn_megamoe_supported(): + self.skipTest("FusedMoeEp is unavailable in the current environment") recipe = MXFP8BlockScaling() if quantization == "mxfp8" else None model, fc1, fc2, _ = self._make_megamoe_model( recipe=recipe, @@ -1561,11 +1629,19 @@ def _run_megamoe_vs_reference( topk_idx, ) + forward_ops = model._module_groups[0]._forward_ops + fused = len(forward_ops) == 1 and isinstance(forward_ops[0][0], FusedMoeEp) + if require_fused: + self.assertTrue(fused, "delay_wgrad_compute test must exercise FusedMoeEp") + if fused: + self.assertTrue(_cudnn_megamoe_supported()) + else: + self.assertFalse(any(isinstance(op, FusedMoeEp) for op, _ in forward_ops)) self.assertEqual(seq_out.dtype, torch.bfloat16) fc1_weight = self._reference_weights(fc1) fc2_weight = self._reference_weights(fc2) - emulate_mxfp8 = quantization == "mxfp8" + emulate_mxfp8 = fused or quantization == "mxfp8" reference = MoeEpReference( num_experts=self.cfg.num_experts, hidden_size=HIDDEN_DIM, @@ -1574,16 +1650,20 @@ def _run_megamoe_vs_reference( ep_group=self.ep_group, max_tokens_per_rank=TOKENS_PER_RANK, output_format=MoeFormat.BF16, - combine_format=MoeFormat.BF16, + combine_format=( + MoeFormat.MXFP8 + if fused and os.environ.get("NVTE_MEGAMOE_MXFP8_COMBINE", "0") == "1" + else MoeFormat.BF16 + ), apply_topk_in_fc1=True, generate_c=True, intermediate_format=MoeFormat.MXFP8 if emulate_mxfp8 else None, - backward_operand_format=MoeFormat.MXFP8 if emulate_mxfp8 else None, + backward_operand_format=MoeFormat.MXFP8 if emulate_mxfp8 and not fused else None, backward_wgrad_mode="operands", token_padding_size=256, weight_interleave_size=32 if quantization == "mxfp8" else None, ) - if emulate_mxfp8: + if emulate_mxfp8 and not fused: reference_activation = quantize_blockwise( tokens.detach(), MoeFormat.MXFP8, @@ -1618,8 +1698,17 @@ def _run_megamoe_vs_reference( ).to(torch.bfloat16) seq_out.backward(dy) if delay_wgrad_compute: + if fused: + self.assertEqual(fc1.wgrad_store.context.qsize(), 1) + self.assertEqual(fc2.wgrad_store.context.qsize(), 1) + if accumulate_into_main_grad: + for op in (fc1, fc2): + self.assertTrue(torch.all(op.weight.main_grad == main_grad_sentinel).item()) fc1.backward_dw() fc2.backward_dw() + if fused: + self.assertTrue(fc1.wgrad_store.context.empty()) + self.assertTrue(fc2.wgrad_store.context.empty()) grad_tokens, grad_topk_weights, wgrad_operands = reference.backward( dy, fc1_weight, @@ -1693,6 +1782,7 @@ def _init_distributed(): runner = unittest.TextTestRunner(stream=sys.stdout, verbosity=2) result = runner.run(suite) dist.barrier() + finalize_moe_ep_resources() ep_finalize() # Deregister symm-mem windows while the comm is still valid. release_symm_mem_pool() diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index cb4a2e76549..d4f65c22222 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -391,6 +391,36 @@ def test_basic_construction_all_same_shape(self) -> None: assert grouped_tensor.get_common_last_dim() == 512 assert grouped_tensor.has_data() + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_module_to_moves_grouped_parameter_storage(self) -> None: + """Module.to preserves GroupedTensor parameters initialized on CPU.""" + grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=2, + shapes=[(4, 8), (4, 8)], + quantizer=None, + device="cpu", + dtype=torch.float32, + ) + values = torch.arange(grouped_tensor.numel(), dtype=torch.float32) + grouped_tensor.rowwise_data.copy_(values) + + module = torch.nn.Module() + module.register_parameter("weight", torch.nn.Parameter(grouped_tensor)) + original_parameter = module.weight + expected = values.to(device="cuda", dtype=torch.bfloat16) + + module.to(device="cuda", dtype=torch.bfloat16) + + assert module.weight is original_parameter + assert isinstance(module.weight, GroupedTensor) + assert module.weight.device.type == "cuda" + assert module.weight.dtype == torch.bfloat16 + assert module.weight.rowwise_data.device.type == "cuda" + assert module.weight.rowwise_data.dtype == torch.bfloat16 + torch.testing.assert_close(module.weight.rowwise_data, expected, rtol=0, atol=0) + members = module.weight.split_into_quantized_tensors() + assert all(member.device.type == "cuda" for member in members) + def test_basic_construction_varying_first_dim(self) -> None: """Test GroupedTensor construction with varying first dimension""" num_tensors = 3 @@ -1710,6 +1740,41 @@ def test_group_dequantize(self, shape: List[Tuple[int, int]]) -> None: for orig, deq in zip(input_tensors, dequantized_tensors): torch.testing.assert_close(deq, orig, atol=0.125, rtol=0.1) + @pytest.mark.parametrize("conversion", ["float", "to"]) + @pytest.mark.parametrize("columnwise", [False, True]) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_grouped_mxfp8_dtype_conversion_for_optimizer( + self, + conversion: str, + columnwise: bool, + ) -> None: + """MXFP8 grouped parameters dequantize before optimizer shard flattening.""" + shapes = [(256, 512), (512, 512)] + inputs = [torch.randn(shape, dtype=torch.bfloat16, device="cuda") for shape in shapes] + quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3) + quantizer.set_usage(rowwise=True, columnwise=columnwise) + grouped = tex.group_quantize( + torch.cat(inputs, dim=0), + quantizer, + len(shapes), + torch.tensor([shape[0] for shape in shapes], dtype=torch.int64, device="cuda"), + ) + + dequantized = grouped.float() if conversion == "float" else grouped.to(dtype=torch.float32) + + assert isinstance(dequantized, GroupedTensor) + assert dequantized.quantizer is None + assert dequantized.dtype == torch.float32 + assert dequantized.num_tensors == grouped.num_tensors + assert dequantized.logical_shape == grouped.logical_shape + assert dequantized.view(-1) is not dequantized + assert dequantized.view(-1).dtype == torch.float32 + assert dequantized.view(-1).numel() == grouped.numel() + + members = dequantized.split_into_quantized_tensors() + for original, member in zip(inputs, members): + torch.testing.assert_close(member, original.float(), atol=0.125, rtol=0.1) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) def test_group_dequantize_cudagraph_capturable(self) -> None: """Ensure group_dequantize is CUDA graph capturable.""" diff --git a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh index fcf8819e5f8..485c11005b8 100644 --- a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh @@ -387,8 +387,6 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output, const bool use_colwise_scaling = input->has_columnwise_data(); NVTE_CHECK(use_rowwise_scaling || use_colwise_scaling, "Input tensor must have either rowwise or columnwise data."); - NVTE_CHECK(!(use_rowwise_scaling && use_colwise_scaling), - "Dequantize only supports rowwise or columnwise scaling, not both simultaneously."); NVTE_CHECK(!input->with_gemm_swizzled_scales, "Input must have scales in compact format."); NVTE_CHECK(!is_fp8_dtype(output->dtype()), "Output must be in higher precision."); @@ -449,6 +447,7 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output, const int64_t *const first_dims_ptr = reinterpret_cast(input->first_dims.dptr); const int64_t *const last_dims_ptr = reinterpret_cast(input->last_dims.dptr); + // Prefer the rowwise representation when both GEMM orientations are available. const e8m0_t *const scales_ptr = use_rowwise_scaling ? reinterpret_cast(input->scale_inv.dptr) : reinterpret_cast(input->columnwise_scale_inv.dptr); diff --git a/transformer_engine/pytorch/ops/fused/__init__.py b/transformer_engine/pytorch/ops/fused/__init__.py index ff1075e6fe2..b11d91f70a9 100644 --- a/transformer_engine/pytorch/ops/fused/__init__.py +++ b/transformer_engine/pytorch/ops/fused/__init__.py @@ -36,3 +36,4 @@ GroupedMLP_CuTeGEMMUnary, GroupedMLP_CuTeGEMMGLUBf16, ) +from .moe_ep import FusedMoeEp # pylint: disable=wrong-import-position diff --git a/transformer_engine/pytorch/ops/fused/moe_ep.py b/transformer_engine/pytorch/ops/fused/moe_ep.py new file mode 100644 index 00000000000..40bf5ae2605 --- /dev/null +++ b/transformer_engine/pytorch/ops/fused/moe_ep.py @@ -0,0 +1,943 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MegaMoE-backed expert-parallel MoE fusion (cudnn.moe_ep.MoeEp).""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +import functools +import os +from typing import Any, Optional + +import torch +import transformer_engine_torch as tex + +from ...constants import MXFP8_BLOCK_SCALING_SIZE +from ...ep import EpConfig +from ...quantization import Recipe +from ...tensor import GroupedTensor, GroupedTensorStorage, Quantizer +from ...utils import mark_grouped_tensor +from .._common import ( + get_accumulate_flag_in_param, + get_dummy_wgrads_for_params, + get_main_grad_from_param, + is_quantized_tensor, + maybe_dequantize, + view_main_grad_as_grouped_buffer, +) +from ..basic import GroupedLinear, MoeCombine, MoeDispatch, ScaledSwiGLU +from ..fuser import register_forward_backward_fusion +from ..op import FusedOperation, FusibleOperation, OperationContext + + +@dataclass +class _MoeEpResource: + """One cached cuDNN operator and its prepared symmetric buffers.""" + + moe: Any + requirements: Any = None + symmetric_buffers: Any = None + device: Optional[torch.device] = None + + +class _MoeEpResourceManager: + """Share one cuDNN MoeEp instance per compatible fused-MoE configuration.""" + + def __init__(self) -> None: + self._resources: dict[tuple[Any, ...], _MoeEpResource] = {} + + def get( + self, + config: EpConfig, + device: torch.device, + intermediate_size: int, + ) -> _MoeEpResource: + """Get or construct the cuDNN operator for one complete configuration.""" + device = torch.device(device) + combine_format = _get_megamoe_combine_format() + key = self._resource_key( + config, + device, + intermediate_size, + combine_format, + ) + resource = self._resources.get(key) + if resource is None: + resource = _MoeEpResource( + moe=self._construct_moe( + config, + intermediate_size, + combine_format, + ), + ) + self._resources[key] = resource + self._prepare(resource, device) + return resource + + @staticmethod + def _construct_moe( + config: EpConfig, + intermediate_size: int, + combine_format: str, + ): + """Construct one cuDNN MoeEp instance from its complete static configuration.""" + moe_ep_api = _import_cudnn_moe_ep() + if moe_ep_api is None: + raise ImportError( + "FusedMoeEp requires cudnn.moe_ep.MoeEp. Install the in-tree " + "cuDNN frontend with: pip install --force-reinstall " + "'./cudnn_frontend[moe_ep]'" + ) + + forward_group_hint = 768 if combine_format == "mxfp8" else 1024 + # DeepSeek-V4-Pro expert shape measured faster with a smaller backward group hint. + is_dsv4_pro_shape = ( + config.hidden_dim == 7168 and intermediate_size == 3072 and config.top_k == 6 + ) + backward_group_hint = 256 if is_dsv4_pro_shape else 512 + resolved_combine_format = ( + moe_ep_api.MoeFormat.MXFP8 if combine_format == "mxfp8" else moe_ep_api.MoeFormat.BF16 + ) + + moe_config = moe_ep_api.MoeEpConfig( + model=moe_ep_api.MoeEpModelConfig( + num_experts=config.num_local_experts * config.ep_group.size(), + hidden_size=config.hidden_dim, + intermediate_size=intermediate_size, + top_k=config.top_k, + ), + parallel=moe_ep_api.MoeEpParallelConfig( + ep_group=config.ep_group, + max_tokens_per_rank=config.max_tokens_per_rank, + physical_recv_pool_rows=config.recv_capacity_per_rank, + drop_on_overflow=config.drop_on_overflow, + token_padding_size=128, + sf_padding_size=128, + ), + data_path=moe_ep_api.MoeEpDataPathConfig( + apply_topk_in_fc1=True, + fc1_weight_layout=(moe_ep_api.MoeEpFc1WeightLayout.GATE_UP_INTERLEAVED_32), + combine_format=resolved_combine_format, + output_format=moe_ep_api.MoeFormat.BF16, + ), + training_forward_tuning=moe_ep_api.MoeEpTuningConfig( + token_back_mode="standalone_warps", + epi_flag_batch=(2, 2), + token_in_flag_batch=8, + group_hint=forward_group_hint, + reduce_topk_in_kernel=False, + ), + training_backward_tuning=moe_ep_api.MoeEpTuningConfig( + token_back_mode="epi_warps", + epi_flag_batch=(4, 2), + token_in_flag_batch=8, + group_hint=backward_group_hint, + reduce_topk_in_kernel=False, + dgrad_optimization="rolling", + ), + training_weight_storage_mode=(moe_ep_api.MoeEpNativeWeightStorageMode.CONTIGUOUS), + validation_mode="trusted", + ) + return moe_ep_api.MoeEp(moe_config) + + @staticmethod + def _resource_key( + config: EpConfig, + device: torch.device, + intermediate_size: int, + combine_format: str, + ) -> tuple[Any, ...]: + """Return every value that affects a shared cuDNN MoeEp instance.""" + return ( + id(config.ep_group), + device.type, + device.index, + config.top_k, + config.hidden_dim, + intermediate_size, + config.num_local_experts, + config.max_tokens_per_rank, + config.recv_capacity_per_rank, + config.alignment, + config.payload_dtype, + config.zero_copy, + config.drop_on_overflow, + combine_format, + ) + + def _prepare( + self, + resource: _MoeEpResource, + device: torch.device, + ) -> tuple[Any, Any]: + """Prepare and cache one resource's symmetric buffer views.""" + device = torch.device(device) + if resource.requirements is None: + resource.requirements = resource.moe.prepare_training(device=device) + resource.symmetric_buffers = resource.moe.training_symmetric_buffers() + resource.device = device + elif resource.device != device: + raise ValueError( + f"shared MoeEp is prepared on {resource.device}, but was requested on {device}" + ) + return resource.requirements, resource.symmetric_buffers + + def cleanup(self) -> None: + """Close every cached MoeEp instance and clear the registry.""" + first_error = None + for resource in tuple(self._resources.values()): + try: + if resource.moe is not None: + resource.moe.close() + except Exception as error: # pylint: disable=broad-exception-caught + if first_error is None: + first_error = error + finally: + resource.requirements = None + resource.symmetric_buffers = None + resource.device = None + resource.moe = None + self._resources.clear() + if first_error is not None: + raise first_error + + +_MOE_EP_RESOURCE_MANAGER = _MoeEpResourceManager() + + +def _cudnn_megamoe_supported() -> bool: + """Whether cuDNN FE provides the stateless training API.""" + if bool(int(os.getenv("NVTE_DISABLE_CUDNN_MEGAMOE", "0"))): + return False + try: + import cudnn + import cudnn.moe_ep as cudnn_moe_ep + except (AttributeError, ImportError): + return False + try: + return all( + hasattr(module, name) + for module, name in ( + (cudnn, "grouped_gemm_wgrad_wrapper_sm100"), + (cudnn_moe_ep, "MoeEp"), + (cudnn_moe_ep, "MoeEpConfig"), + (cudnn_moe_ep, "MoeEpModelConfig"), + (cudnn_moe_ep, "MoeEpParallelConfig"), + (cudnn_moe_ep, "MoeEpDataPathConfig"), + (cudnn_moe_ep, "MoeEpFc1WeightLayout"), + (cudnn_moe_ep, "MoeEpNativeWeightStorageMode"), + (cudnn_moe_ep, "MoeEpTuningConfig"), + (cudnn_moe_ep, "MoeFormat"), + (cudnn_moe_ep.MoeEp, "training_symmetric_buffers"), + (cudnn_moe_ep, "MoeEpNativeForwardWeights"), + (cudnn_moe_ep, "MoeEpNativeBackwardWeights"), + (cudnn_moe_ep, "MoeEpTrainingForwardOutputs"), + (cudnn_moe_ep, "MoeEpTrainingBackwardOutputs"), + (cudnn_moe_ep, "MoeEpTrainingWgradOperands"), + ) + ) + except (AttributeError, ImportError): + return False + + +def finalize_moe_ep_resources() -> None: + """Close all cached cuDNN/NVSHMEM resources before EP group teardown.""" + # CUDA graph callables form reference cycles. Collect unreachable graph + # executables before tearing down communication resources they captured. + import gc + + gc.collect() + devices = { + resource.device + for resource in _MOE_EP_RESOURCE_MANAGER._resources.values() + if resource.device is not None + } + for device in devices: + torch.cuda.synchronize(device) + _MOE_EP_RESOURCE_MANAGER.cleanup() + + +def _get_megamoe_combine_format() -> str: + """Return the MegaMoE combine wire format selected by the environment.""" + enabled = int(os.environ.get("NVTE_MEGAMOE_MXFP8_COMBINE", "0")) + return "mxfp8" if enabled > 0 else "bf16" + + +def _allocate_training_buffer(requirements, name: str, device: torch.device) -> torch.Tensor: + """Allocate one caller-owned buffer from cuDNN's named contract.""" + shape, stride, dtype, _alignment = requirements[name] + return torch.empty_strided(shape, stride, dtype=dtype, device=device) + + +def _quantize_into_cudnn_symmetric_buffer( + input_: torch.Tensor, + data_buffer: torch.Tensor, + scale_buffer: torch.Tensor, + block_scaled_cls: type, +): + """Quantize a plain tensor directly into cuDNN's symmetric MXFP8 storage.""" + from ...tensor.mxfp8_tensor import MXFP8Quantizer + from ...tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage + + token_count, hidden_size = input_.shape + scale_rows = (token_count + 127) // 128 * 128 + quantizer = MXFP8Quantizer( + tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + ) + data = data_buffer[:token_count] + padded_scale = scale_buffer[:scale_rows] + output = MXFP8TensorStorage( + data.view(torch.uint8), + padded_scale.view(torch.uint8), + None, + None, + tex.DType.kFloat8E4M3, + quantizer, + False, + fake_dtype=input_.dtype, + ) + tex.quantize(input_, quantizer, output, None) + return block_scaled_cls( + data=data, + scale=padded_scale[:token_count, : hidden_size // MXFP8_BLOCK_SCALING_SIZE], + format="mxfp8", + logical_shape=tuple(input_.shape), + axis=1, + ) + + +def _launch_grouped_wgrad_from_operands( + x_operands: list[torch.Tensor], + dy_operands: list[torch.Tensor], + output: torch.Tensor | GroupedTensor | list[torch.Tensor], + *, + packed_output: Optional[torch.Tensor] = None, + offsets: torch.Tensor, + accumulate: bool, + descriptor_workspace: torch.Tensor, +) -> None: + """Compute one TE-layout grouped wgrad directly from MegaMoE's operands.""" + from cudnn import grouped_gemm_wgrad_wrapper_sm100 + + x, x_scale = x_operands + dy, dy_scale = dy_operands + output = packed_output if packed_output is not None else output + output_data = ( + output.rowwise_data.view(output.shape) if isinstance(output, GroupedTensor) else output + ) + # Present the same dY^T @ X convention used by grouped MLP. MegaMoE exports + # both layers' activations directly in contiguous token-major form. + # The public wrapper selects the Rubin specialization on SM107. + grouped_gemm_wgrad_wrapper_sm100( + a_tensor=dy.transpose(0, 1), + b_tensor=x, + sfa_tensor=dy_scale, + sfb_tensor=x_scale, + offsets_tensor=offsets, + output_mode="dense", + wgrad_tensor=output_data, + wgrad_dtype=output_data.dtype, + acc_dtype=torch.float32, + sf_vec_size=MXFP8_BLOCK_SCALING_SIZE, + accumulate_on_output=accumulate, + input_order="tensor2d", + descriptor_workspace=descriptor_workspace, + ) + + +def _compute_grouped_weight_grad( + op: GroupedLinear, + operands, + prefix: str, + descriptor_workspace: torch.Tensor, +) -> list[Optional[torch.Tensor]]: + """Compute one dense weight gradient with cuDNN's grouped-WGrad API.""" + weight = op.weight + if not weight.requires_grad: + return [None] + + weight_shape = (op.out_features, op.in_features) + output_shape = (op.num_groups, *weight_shape) + accumulate = False + if op._accumulate_into_main_grad: + output_data = get_main_grad_from_param( + weight, + op_label=f"FusedMoeEp {prefix.upper()}", + ) + output_data = view_main_grad_as_grouped_buffer( + output_data, + op.num_groups, + weight_shape, + label=f"FusedMoeEp {prefix.upper()} weight", + ) + accumulate = get_accumulate_flag_in_param(weight) + else: + output_data = torch.empty( + output_shape, + dtype=weight.dtype, + device=weight.device, + ) + + x_operands = [ + getattr(operands, f"{prefix}_a"), + getattr(operands, f"{prefix}_sfa"), + ] + dy_operands = [ + getattr(operands, f"{prefix}_b"), + getattr(operands, f"{prefix}_sfb"), + ] + if op.wgrad_store.delay_wgrad_compute(): + grad_views = [output_data[i] for i in range(op.num_groups)] + wgrad_fn = functools.partial( + _launch_grouped_wgrad_from_operands, + packed_output=output_data, + offsets=operands.expert_offsets, + accumulate=accumulate, + descriptor_workspace=descriptor_workspace, + ) + op.wgrad_store.put([x_operands, dy_operands, grad_views], wgrad_fn) + if op._accumulate_into_main_grad: + return get_dummy_wgrads_for_params([weight]) + return [None] + + _launch_grouped_wgrad_from_operands( + x_operands, + dy_operands, + output_data, + offsets=operands.expert_offsets, + accumulate=accumulate, + descriptor_workspace=descriptor_workspace, + ) + + if op._accumulate_into_main_grad: + return get_dummy_wgrads_for_params([weight]) + return [output_data] + + +def _grouped_linear_supported(op: GroupedLinear) -> bool: + weight = op.weight if op.single_grouped_weight else None + weight_ok = ( + isinstance(weight, GroupedTensor) + and weight.dtype is torch.bfloat16 + and weight.rowwise_data is not None + ) + if weight_ok and weight.quantizer is not None: + recipe = weight.quantizer._get_compatible_recipe() + weight_ok = ( + recipe is not None + and recipe.mxfp8() + and not weight._with_gemm_swizzled_scales + and weight.rowwise_data is not None + and weight.scale_inv is not None + and weight.columnwise_data is not None + and weight.columnwise_scale_inv is not None + ) + else: + weight_ok = False + + return ( + not op.use_bias + and not op._scale_bias + and op.single_grouped_weight + and not op.single_grouped_bias + and not op._is_distributed_weight() + and weight_ok + ) + + +def _import_cudnn_moe_ep(): + """Return the ``cudnn.moe_ep`` API module or ``None`` if it is missing.""" + try: + import cudnn.moe_ep as moe_ep_api + except ImportError: + return None + return moe_ep_api + + +def _routing_extras_internal( + dispatch: MoeDispatch, + fc1: GroupedLinear, + activation: ScaledSwiGLU, + fc2: GroupedLinear, +) -> bool: + """Whether the dispatch routing extras stay inside the fusion. + + The fused op keeps tokens-per-expert and the received routing weights + internal to :class:`cudnn.moe_ep.MoeEp`, so it can only replace the + sequence when those two outputs feed exactly these ops and are not + returned to the caller. + """ + tokens_per_expert, routing_weights = dispatch._extra_output_channels + if tokens_per_expert is None or routing_weights is None: + return False + if any(dispatch._extra_output_to_caller): + return False + return ( + fc1._extra_input_channels[0] == tokens_per_expert + and fc2._extra_input_channels[0] == tokens_per_expert + and activation._extra_input_channels[0] == routing_weights + ) + + +def _is_mxfp8_comms_recipe(recipe: Optional[Recipe]) -> bool: + return recipe is not None and recipe.mxfp8() + + +def _megamoe_supported(config: EpConfig, fc2: GroupedLinear) -> bool: + """Static MegaMoE capability gates that can be checked before first launch.""" + # MegaMoE always moves dispatched tokens and combine gradients as MXFP8. + if not ( + _is_mxfp8_comms_recipe(config.dispatch_fwd_quant_recipe) + and _is_mxfp8_comms_recipe(config.combine_bwd_quant_recipe) + ): + return False + if not _cudnn_megamoe_supported(): + return False + if _import_cudnn_moe_ep() is None: + return False + if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7): + return False + if config.max_tokens_per_rank <= 0: + return False + if config.recv_capacity_per_rank is None or config.recv_capacity_per_rank <= 0: + return False + if config.hidden_dim % 128 != 0 or fc2.in_features % 256 != 0: + return False + if config.top_k > 32: + return False + return True + + +def is_moe_fusion_supported( + window: Sequence[FusibleOperation], + recipe: Optional[Recipe], +) -> bool: + """Whether a five-op Sequential window supports FusedMoeEp.""" + if len(window) != 5: + return False + if recipe is not None and not recipe.mxfp8(): + return False + dispatch, fc1, activation, fc2, combine = window + if not ( + isinstance(dispatch, MoeDispatch) + and isinstance(fc1, GroupedLinear) + and isinstance(activation, ScaledSwiGLU) + and isinstance(fc2, GroupedLinear) + and isinstance(combine, MoeCombine) + ): + return False + config = dispatch.config + if combine.config != config or config.payload_dtype is not torch.bfloat16: + return False + if not (_grouped_linear_supported(fc1) and _grouped_linear_supported(fc2)): + return False + if activation.activation_recompute_in_mlp or activation.glu_interleave_size != 32: + return False + if not _routing_extras_internal(dispatch, fc1, activation, fc2): + return False + if not _megamoe_supported(config, fc2): + return False + return ( + fc1.num_groups == config.num_local_experts + and fc2.num_groups == config.num_local_experts + and fc1.in_features == config.hidden_dim + and fc2.out_features == config.hidden_dim + and fc1.out_features == 2 * fc2.in_features + ) + + +class FusedMoeEp(FusedOperation): + """Joint EP MoE fusion implemented with :class:`cudnn.moe_ep.MoeEp`.""" + + def __init__( + self, + *, + dispatch: MoeDispatch, + fc1: GroupedLinear, + activation: ScaledSwiGLU, + fc2: GroupedLinear, + combine: MoeCombine, + ) -> None: + super().__init__([dispatch, fc1, activation, fc2, combine]) + from cudnn.moe_ep import BlockScaledTensor + + self._block_scaled_cls = BlockScaledTensor + self._resource = None + + def _prepare_training(self, device: torch.device) -> None: + """Acquire and prepare the shared cuDNN runtime on the execution device.""" + if self._resource is not None: + return + resource = _MOE_EP_RESOURCE_MANAGER.get( + self.dispatch.config, + device, + self.fc2.in_features, + ) + self._resource = resource + + def _make_native_training_weights(self): + """Expose TE payloads with freshly swizzled caller-owned scale buffers.""" + from cudnn.moe_ep import ( + MoeEpNativeBackwardWeights, + MoeEpNativeForwardWeights, + MoeEpNativeWeight, + MoeEpNativeWeightLayout, + ) + + def swizzle(weight: GroupedTensor): + native = weight.copy() + tex.grouped_swizzle_for_gemm(native, rowwise=True, columnwise=True) + return native + + fc1 = swizzle(self.fc1.weight) + fc2 = swizzle(self.fc2.weight) + num_experts = self.fc1.num_groups + + def data(tensor: torch.Tensor, shape: tuple[int, ...]) -> torch.Tensor: + return tensor.view(*shape).view(torch.float8_e4m3fn) + + def scale(tensor: torch.Tensor) -> torch.Tensor: + return tensor.view(torch.float8_e8m0fnu).reshape(num_experts, -1) + + fc1_rowwise = data( + fc1.rowwise_data, + (num_experts, self.fc1.out_features, self.fc1.in_features), + ).permute(0, 2, 1) + fc2_rowwise = data( + fc2.rowwise_data, + (num_experts, self.fc2.out_features, self.fc2.in_features), + ).permute(0, 2, 1) + fc1_columnwise = data( + fc1.columnwise_data, + (num_experts, self.fc1.out_features, self.fc1.in_features), + ) + fc2_columnwise = data( + fc2.columnwise_data, + (num_experts, self.fc2.out_features, self.fc2.in_features), + ) + forward = MoeEpNativeForwardWeights( + fc1=MoeEpNativeWeight( + fc1_rowwise, + scale(fc1.scale_inv), + MoeEpNativeWeightLayout.FORWARD_FC1_GATE_UP_INTERLEAVED_32_V1, + ), + fc2=MoeEpNativeWeight( + fc2_rowwise, + scale(fc2.scale_inv), + MoeEpNativeWeightLayout.FORWARD_FC2_K_MAJOR_V1, + ), + ) + backward = MoeEpNativeBackwardWeights( + w2_transpose=MoeEpNativeWeight( + fc2_columnwise, + scale(fc2.columnwise_scale_inv), + MoeEpNativeWeightLayout.BACKWARD_W2_TRANSPOSE_V1, + ), + w1_transpose=MoeEpNativeWeight( + fc1_columnwise, + scale(fc1.columnwise_scale_inv), + MoeEpNativeWeightLayout.BACKWARD_W1_TRANSPOSE_GATE_UP_INTERLEAVED_32_V1, + ), + ) + return forward, backward + + def _make_training_forward_outputs(self, device: torch.device): + """Allocate forward output and backward-state destinations.""" + from cudnn.moe_ep import MoeEpTrainingForwardOutputs + + requirements = self._resource.requirements + return MoeEpTrainingForwardOutputs( + output=self._resource.symmetric_buffers["output"], + fc1_preact=_allocate_training_buffer(requirements, "fc1_preact", device), + fc1_a=_allocate_training_buffer(requirements, "fc1_a", device), + fc1_sfa=_allocate_training_buffer(requirements, "fc1_sfa", device), + valid_route_counts=_allocate_training_buffer( + requirements, + "valid_route_counts", + device, + ), + expert_offsets=_allocate_training_buffer(requirements, "expert_offsets", device), + ) + + def _make_training_backward_outputs(self, device: torch.device): + """Allocate gradient and grouped-WGrad operand destinations.""" + from cudnn.moe_ep import MoeEpTrainingBackwardOutputs + + requirements = self._resource.requirements + return MoeEpTrainingBackwardOutputs( + grad_activation=self._resource.symmetric_buffers["grad_activation"], + dprob=self._resource.symmetric_buffers["dprob"], + fc1_b=_allocate_training_buffer(requirements, "fc1_b", device), + fc1_sfb=_allocate_training_buffer(requirements, "fc1_sfb", device), + fc2_a=_allocate_training_buffer(requirements, "fc2_a", device), + fc2_sfa=_allocate_training_buffer(requirements, "fc2_sfa", device), + fc2_b=_allocate_training_buffer(requirements, "fc2_b", device), + fc2_sfb=_allocate_training_buffer(requirements, "fc2_sfb", device), + ) + + def _make_training_wgrad_workspaces(self): + """Return caller-owned descriptor workspaces for both FC gradients. + + Each GroupedLinear owns a persistent workspace. Wgrad kernels without + tensormap proxy fences can read stale cached TMA descriptors when a + freshly allocated workspace reuses an address with new contents. + """ + from cudnn import get_grouped_gemm_wgrad_workspace_size_sm100 + + workspace_bytes = get_grouped_gemm_wgrad_workspace_size_sm100( + self.fc1.num_groups, + output_mode="dense", + input_order="tensor2d", + ) + workspaces = [] + for fc_op in (self.fc1, self.fc2): + device = fc_op.weight.device + workspace = getattr(fc_op, "_cudnn_wgrad_workspace", None) + if ( + workspace is None + or workspace.numel() != workspace_bytes + or workspace.device != device + ): + workspace = torch.empty(workspace_bytes, dtype=torch.uint8, device=device) + fc_op._cudnn_wgrad_workspace = workspace + workspaces.append(workspace) + return tuple(workspaces) + + @property + def dispatch(self) -> MoeDispatch: + """Return the underlying dispatch operation.""" + return self.basic_ops[0] + + @property + def fc1(self) -> GroupedLinear: + """Return the first grouped linear operation.""" + return self.basic_ops[1] + + @property + def fc2(self) -> GroupedLinear: + """Return the second grouped linear operation.""" + return self.basic_ops[3] + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, Sequence[Sequence[Optional[torch.Tensor]]]]: + if ( + not isinstance(input_, torch.Tensor) + or is_quantized_tensor(input_) + or input_.dtype is not torch.bfloat16 + ): + raise TypeError( + f"FusedMoeEp input must be a plain BF16 tensor, got {type(input_).__name__}." + ) + + topk_idx, topk_weights = basic_op_extra_inputs[0] + if topk_weights.dtype is not torch.float32: + raise TypeError(f"topk_weights must be float32, got {topk_weights.dtype}.") + activation = input_ + if not activation.is_contiguous(): + activation = activation.contiguous() + if topk_idx.dtype is not torch.int32: + topk_idx = topk_idx.to(dtype=torch.int32) + if not topk_idx.is_contiguous(): + topk_idx = topk_idx.contiguous() + if not topk_weights.is_contiguous(): + topk_weights = topk_weights.contiguous() + + self._prepare_training(activation.device) + activation = _quantize_into_cudnn_symmetric_buffer( + activation, + self._resource.symmetric_buffers["forward_input"], + self._resource.symmetric_buffers["forward_input_scale"], + self._block_scaled_cls, + ) + forward_weights, backward_weights = self._make_native_training_weights() + forward_out = self._make_training_forward_outputs(activation.device) + output = self._resource.moe.training_forward( + activation, + topk_idx, + topk_weights, + weights=forward_weights, + out=forward_out, + ) + + if any(ctx.requires_grad for ctx in basic_op_ctxs): + fc1_a = GroupedTensorStorage( + shape=( + forward_out.fc1_a.numel() // self.fc1.in_features, + self.fc1.in_features, + ), + dtype=forward_out.fc1_a.dtype, + num_tensors=self.fc1.num_groups, + columnwise_data=forward_out.fc1_a, + columnwise_scale_inv=forward_out.fc1_sfa, + ) + mark_grouped_tensor(forward_out.fc1_preact, fc1_a) + active_rows = forward_out.expert_offsets[-1:] + forward_out.fc1_preact.grouped_tensor_num_tokens = active_rows + forward_out.fc1_a.grouped_tensor_num_tokens = active_rows + forward_out.fc1_sfa.grouped_tensor_num_tokens = active_rows + basic_op_ctxs[0].save_for_backward( + topk_idx, + topk_weights, + forward_out.fc1_preact, + forward_out.fc1_a, + forward_out.fc1_sfa, + forward_out.valid_route_counts, + forward_out.expert_offsets, + backward_weights.w2_transpose.payload, + backward_weights.w2_transpose.scale, + backward_weights.w1_transpose.payload, + backward_weights.w1_transpose.scale, + ) + + return output, [ + (None, None), + (), + (), + (), + (), + ] + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + *, + basic_op_grad_extra_outputs: list[tuple[Optional[torch.Tensor], ...]], + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + del basic_op_grad_extra_outputs + grad_output = maybe_dequantize( + grad_output, + torch.bfloat16, + ) + if not grad_output.is_contiguous(): + grad_output = grad_output.contiguous() + grad_output = _quantize_into_cudnn_symmetric_buffer( + grad_output, + self._resource.symmetric_buffers["backward_input"], + self._resource.symmetric_buffers["backward_input_scale"], + self._block_scaled_cls, + ) + ( + topk_idx, + topk_weights, + fc1_preact, + fc1_a, + fc1_sfa, + valid_route_counts, + expert_offsets, + w2_payload, + w2_scale, + w1_payload, + w1_scale, + ) = basic_op_ctxs[0].saved_tensors + from cudnn.moe_ep import ( + MoeEpNativeBackwardWeights, + MoeEpNativeWeight, + MoeEpNativeWeightLayout, + ) + + backward_weights = MoeEpNativeBackwardWeights( + w2_transpose=MoeEpNativeWeight( + w2_payload, + w2_scale, + MoeEpNativeWeightLayout.BACKWARD_W2_TRANSPOSE_V1, + ), + w1_transpose=MoeEpNativeWeight( + w1_payload, + w1_scale, + MoeEpNativeWeightLayout.BACKWARD_W1_TRANSPOSE_GATE_UP_INTERLEAVED_32_V1, + ), + ) + backward_out = self._make_training_backward_outputs(grad_output.device) + grad_input, grad_topk_weights, wgrad_operands = self._resource.moe.training_backward( + grad_output, + topk_idx, + topk_weights, + weights=backward_weights, + fc1_preact=fc1_preact, + fc1_a=fc1_a, + fc1_sfa=fc1_sfa, + valid_route_counts=valid_route_counts, + expert_offsets=expert_offsets, + out=backward_out, + ) + fc1_workspace, fc2_workspace = self._make_training_wgrad_workspaces() + fc1_param_grads = _compute_grouped_weight_grad( + self.fc1, + wgrad_operands, + "fc1", + fc1_workspace, + ) + fc2_param_grads = _compute_grouped_weight_grad( + self.fc2, + wgrad_operands, + "fc2", + fc2_workspace, + ) + return ( + grad_input, + [(), fc1_param_grads, (), fc2_param_grads, ()], + [(None, grad_topk_weights.float()), (None,), (None,), (None,), (None,)], + ) + + +def fuse_ops( + ops: list[FusibleOperation], + *, + recipe: Optional[Recipe] = None, + **unused: Any, +) -> list[FusibleOperation]: + """Fuse supported five-op EP MoE sequences into MegaMoE. + + Unfused Sequential (NCCL dispatch/combine + GroupedLinear + ScaledSwiGLU) + is the default. MegaMoE is claimed only when + :func:`is_moe_fusion_supported` succeeds. + """ + del unused + out: list[FusibleOperation] = [] + idx = 0 + while idx < len(ops): + window = ops[idx : idx + 5] + if is_moe_fusion_supported(window, recipe): + dispatch, fc1, activation, fc2, combine = window + out.append( + FusedMoeEp( + dispatch=dispatch, + fc1=fc1, + activation=activation, + fc2=fc2, + combine=combine, + ) + ) + idx += 5 + else: + out.append(ops[idx]) + idx += 1 + return out + + +register_forward_backward_fusion(fuse_ops, prepend=True) + + +__all__ = [ + "FusedMoeEp", + "finalize_moe_ep_resources", + "is_moe_fusion_supported", +] diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index 3aa1a21737e..12adb01409a 100644 --- a/transformer_engine/pytorch/tensor/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/grouped_tensor.py @@ -8,8 +8,10 @@ from typing import List, Optional, Tuple import torch +import transformer_engine_torch as tex from torch.utils._pytree import tree_map +from ..constants import TE_DType from transformer_engine.pytorch.device import te_device_type from ..quantized_tensor import QuantizedTensorStorage, Quantizer from .storage.grouped_tensor_storage import GroupedTensorStorage @@ -177,6 +179,119 @@ def __new__( ) return instance + def __tensor_flatten__(self): + """Expose grouped backing buffers to PyTorch's wrapper-subclass protocol.""" + tensor_attrs = ( + "rowwise_data", + "columnwise_data", + "scale_inv", + "columnwise_scale_inv", + "amax", + "columnwise_amax", + "scale", + "first_dims", + "last_dims", + "tensor_offsets", + ) + present = [name for name in tensor_attrs if getattr(self, name) is not None] + context = { + "cls": type(self), + "logical_shape": self.logical_shape, + "fake_dtype": self.fake_dtype, + "num_tensors": self.num_tensors, + "tensor_shapes": self.tensor_shapes, + "quantizer": self.quantizer, + "offsets": self.offsets, + "scale_inv_offsets": self.scale_inv_offsets, + "columnwise_scale_inv_offsets": self.columnwise_scale_inv_offsets, + "requires_grad": self.requires_grad, + "with_gemm_swizzled_scales": self._with_gemm_swizzled_scales, + "row_scaled_nvfp4": self.row_scaled_nvfp4, + "nvfp4_use_4over6": self.nvfp4_use_4over6, + "nvfp4_e4m3_max": self.nvfp4_e4m3_max, + } + return present, context + + def dequantize(self, dtype: Optional[torch.dtype] = None) -> GroupedTensor: + """Return a high-precision grouped tensor with the same logical members.""" + if dtype is None: + dtype = self.get_dtype() + if dtype not in TE_DType: + raise TypeError(f"Unsupported dequantization dtype: {dtype}") + + if self.quantizer is None: + if dtype == self.dtype: + return self + return self.to(dtype=dtype) + + dequantized = tex.group_dequantize(self, TE_DType[dtype]) + out = type(self)( + shape=self.logical_shape, + dtype=dtype, + num_tensors=self.num_tensors, + shapes=self.tensor_shapes, + quantizer=None, + data=dequantized.rowwise_data, + columnwise_data=None, + scale_inv=None, + columnwise_scale_inv=None, + amax=None, + columnwise_amax=None, + scale=None, + first_dims=self.first_dims, + last_dims=self.last_dims, + tensor_offsets=self.tensor_offsets, + offsets=self.offsets, + scale_inv_offsets=None, + columnwise_scale_inv_offsets=None, + requires_grad=False, + with_gemm_swizzled_scales=False, + ) + return out + + def float(self) -> GroupedTensor: + """Dequantize grouped storage to FP32.""" + return self.dequantize(dtype=torch.float32) + + def bfloat16(self) -> GroupedTensor: + """Dequantize grouped storage to BF16.""" + return self.dequantize(dtype=torch.bfloat16) + + def half(self) -> GroupedTensor: + """Dequantize grouped storage to FP16.""" + return self.dequantize(dtype=torch.float16) + + @staticmethod + def __tensor_unflatten__(inner_tensors, context, outer_size, outer_stride): + """Rebuild a GroupedTensor from PyTorch-managed backing buffers.""" + del outer_size + return context["cls"]( + shape=context["logical_shape"], + dtype=context["fake_dtype"], + num_tensors=context["num_tensors"], + shapes=context["tensor_shapes"], + quantizer=context["quantizer"], + data=inner_tensors.get("rowwise_data"), + columnwise_data=inner_tensors.get("columnwise_data"), + scale_inv=inner_tensors.get("scale_inv"), + columnwise_scale_inv=inner_tensors.get("columnwise_scale_inv"), + amax=inner_tensors.get("amax"), + columnwise_amax=inner_tensors.get("columnwise_amax"), + scale=inner_tensors.get("scale"), + first_dims=inner_tensors.get("first_dims"), + last_dims=inner_tensors.get("last_dims"), + tensor_offsets=inner_tensors.get("tensor_offsets"), + offsets=context["offsets"], + scale_inv_offsets=context["scale_inv_offsets"], + columnwise_scale_inv_offsets=context["columnwise_scale_inv_offsets"], + requires_grad=context["requires_grad"], + stride=list(outer_stride) if outer_stride is not None else None, + with_gemm_swizzled_scales=context["with_gemm_swizzled_scales"], + row_scaled_nvfp4=context["row_scaled_nvfp4"], + nvfp4_use_4over6=context["nvfp4_use_4over6"], + nvfp4_e4m3_max=context["nvfp4_e4m3_max"], + ) + @classmethod def __torch_dispatch__(cls, func, types, args, kwargs=None): """Dispatch by dequantizing grouped members, then requantizing writes.""" @@ -245,6 +360,74 @@ def make_wrapper_like(src: GroupedTensor) -> GroupedTensor: tensor.record_stream(stream) return None + # Module.to()/cuda() use aten._to_copy on parameters. Rebuild the + # wrapper around migrated grouped storage instead of taking the default + # dequantize-and-stack path, which loses the GroupedTensor subclass and + # leaves its backing buffers on the original device. + if func == torch.ops.aten._to_copy.default: + src = args[0] + if not isinstance(src, GroupedTensor): + raise TypeError(f"Expected GroupedTensor, got {type(src).__name__}") + + target_device = kwargs.get("device", src.device) + target_dtype = kwargs.get("dtype", src.dtype) + target_layout = kwargs.get("layout", src.layout) + pin_memory = kwargs.get("pin_memory", False) + non_blocking = kwargs.get("non_blocking", False) + + if target_layout != torch.strided: + raise NotImplementedError( + f"{cls.__name__} only supports strided layout, got {target_layout}" + ) + if pin_memory: + raise NotImplementedError(f"{cls.__name__} does not support pin_memory=True") + if src.quantizer is not None and target_dtype != src.dtype: + out = src.dequantize(dtype=target_dtype) + if target_device != src.device: + out = out.to(device=target_device, non_blocking=non_blocking) + return out + + def move_storage( + tensor: Optional[torch.Tensor], *, convert_dtype: bool = False + ) -> Optional[torch.Tensor]: + if tensor is None: + return None + dtype = target_dtype if convert_dtype else tensor.dtype + return tensor.to( + device=target_device, + dtype=dtype, + non_blocking=non_blocking, + copy=True, + ) + + convert_data_dtype = src.quantizer is None + return type(src)( + shape=src.logical_shape, + dtype=target_dtype, + num_tensors=src.num_tensors, + shapes=src.tensor_shapes, + quantizer=src.quantizer, + data=move_storage(src.rowwise_data, convert_dtype=convert_data_dtype), + columnwise_data=move_storage(src.columnwise_data, convert_dtype=convert_data_dtype), + scale_inv=move_storage(src.scale_inv), + columnwise_scale_inv=move_storage(src.columnwise_scale_inv), + amax=move_storage(src.amax), + columnwise_amax=move_storage(src.columnwise_amax), + scale=move_storage(src.scale), + first_dims=move_storage(src.first_dims), + last_dims=move_storage(src.last_dims), + tensor_offsets=move_storage(src.tensor_offsets), + offsets=src.offsets, + scale_inv_offsets=src.scale_inv_offsets, + columnwise_scale_inv_offsets=src.columnwise_scale_inv_offsets, + requires_grad=src.requires_grad, + stride=list(src.stride()), + with_gemm_swizzled_scales=src._with_gemm_swizzled_scales, + row_scaled_nvfp4=src.row_scaled_nvfp4, + nvfp4_use_4over6=src.nvfp4_use_4over6, + nvfp4_e4m3_max=src.nvfp4_e4m3_max, + ) + # Parameter construction may invoke aten.expand on tensor subclasses. # Handle this explicitly so grouped parameters can be created safely. if func == torch.ops.aten.expand.default: