diff --git a/megatron/core/inference/moe/activations.py b/megatron/core/inference/moe/activations.py index 929fbb7b9e4..4330388ddb5 100644 --- a/megatron/core/inference/moe/activations.py +++ b/megatron/core/inference/moe/activations.py @@ -318,4 +318,6 @@ def squared_relu_and_quantize_mxfp8( NUM_BLOCKS=NUM_BLOCKS, ) - return MXFP8Tensor(data=out_fp8, scale=out_scale.view(torch.float8_e8m0fnu), backend="triton") + return MXFP8Tensor( + data=out_fp8, scale=out_scale.view(torch.float8_e8m0fnu), dtype=x.dtype, backend="triton" + ) diff --git a/megatron/core/inference/moe/permute.py b/megatron/core/inference/moe/permute.py index f7785b49edf..15c228f1171 100644 --- a/megatron/core/inference/moe/permute.py +++ b/megatron/core/inference/moe/permute.py @@ -728,6 +728,9 @@ def permute_and_quantize_mxfp8( ) permuted_mxfp8 = MXFP8Tensor( - data=out_fp8, scale=out_scale.view(torch.float8_e8m0fnu), backend="triton" + data=out_fp8, + scale=out_scale.view(torch.float8_e8m0fnu), + dtype=hidden_states.dtype, + backend="triton", ) return permuted_mxfp8, permuted_probs, permutation_map, inclusive_expert_offsets diff --git a/megatron/core/inference/quantization/mxfp8_tensor.py b/megatron/core/inference/quantization/mxfp8_tensor.py index 702fa4a8ed7..8276cbfcb0d 100644 --- a/megatron/core/inference/quantization/mxfp8_tensor.py +++ b/megatron/core/inference/quantization/mxfp8_tensor.py @@ -1,6 +1,6 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Literal, Optional import torch @@ -122,6 +122,22 @@ class MXFP8Tensor: data: torch.Tensor # [M, K] fp8_e4m3fn scale: torch.Tensor # 1D swizzled or [M, K // 32] unswizzled scales backend: Optional[MXFP8Backend] = None # quantization and GEMM backend + # Keyword-only preserves the pre-existing third positional ``backend`` argument. + # This metadata supports tensor-like conversion consumers such as Megatron + # Bridge; it does not constrain the precision accepted by ``quantize_``/``copy_``. + # Legacy direct constructors have unknown logical dtype until their next + # successful update. + dtype: Optional[torch.dtype] = field(default=None, kw_only=True) + + @property + def shape(self) -> torch.Size: + """Shape of the quantized data storage.""" + return self.data.shape + + @property + def device(self) -> torch.device: + """Device holding the quantized tensor.""" + return self.data.device def size(self, idx: Optional[int] = None): """Wrapper for calling self.data.size()""" @@ -141,12 +157,60 @@ def scale_2d(self, K: Optional[int] = None) -> torch.Tensor: padded_cols = n_col_blocks * MXFP8_SCALE_COL_BLOCK return self.scale.reshape(-1, padded_cols) + def quantize_(self, value: torch.Tensor) -> "MXFP8Tensor": + """Quantize a logical tensor into existing storage without changing pointers. + + The source dtype is preserved when supported by the configured backend so + the quantizer does not introduce an avoidable intermediate downcast. + FlashInfer FP32 inputs are converted to BF16 because that backend accepts + FP16/BF16 inputs only. For legacy instances with unknown logical dtype, the + first successful update records the dtype actually passed to the quantizer. + Shape broadcasting is unsupported because MXFP8 scale storage has fixed + geometry. + """ + if self.data.ndim != 2: + raise ValueError( + "In-place MXFP8 updates require 2D destination storage; " + f"got shape {tuple(self.data.shape)}." + ) + if value.shape != self.shape: + raise ValueError( + f"MXFP8 shape mismatch: expected {tuple(self.shape)}, got {tuple(value.shape)}." + ) + if self.backend is None: + raise ValueError("Cannot update an MXFP8Tensor without a quantization backend.") + value = value.to(device=self.device) + if self.backend == "flashinfer" and value.dtype == torch.float32: + value = value.to(dtype=torch.bfloat16) + value = value.contiguous() + quantized = MXFP8Tensor.from_bf16(value, backend=self.backend) + self.data.copy_(quantized.data) + self.scale.view(torch.uint8).copy_(quantized.scale.view(torch.uint8)) + if self.dtype is None: + self.dtype = value.dtype + return self + + def copy_(self, value: torch.Tensor) -> "MXFP8Tensor": + """Tensor-compatible alias for :meth:`quantize_` used by generic writeback. + + Unlike ``torch.Tensor.copy_``, this method requires an exact shape and + does not accept ``non_blocking``; quantization and storage updates are + ordered on the current CUDA stream. The alias lets external conversion + integrations such as Megatron Bridge treat plain and MXFP8 destinations + uniformly while ``quantize_`` remains available to existing callers. + """ + return self.quantize_(value) + @classmethod def from_bf16(cls, x: torch.Tensor, group_size: int = 32, backend: MXFP8Backend = "flashinfer"): - """Quantize BF16 tensor to MXFP8. + """Quantize a floating-point CUDA tensor to MXFP8. + + The historical method name is retained for compatibility. The Triton + backend accepts BF16, FP16, and FP32 inputs; FlashInfer accepts BF16 and + FP16 inputs. Args: - x: [M, K] BF16 tensor on CUDA. + x: [M, K] floating-point tensor on CUDA. group_size: MXFP8 group size (default 32). backend: 'triton' (fused quantize + swizzle Triton kernel) or 'flashinfer' (single fused FlashInfer CUDA kernel). @@ -168,4 +232,4 @@ def from_bf16(cls, x: torch.Tensor, group_size: int = 32, backend: MXFP8Backend f"Unknown MXFP8 quantization backend: '{backend}'. " "Must be 'triton' or 'flashinfer'." ) - return cls(data=data, scale=scale, backend=backend) + return cls(data=data, scale=scale, backend=backend, dtype=x.dtype) diff --git a/megatron/core/inference/quantization/utils.py b/megatron/core/inference/quantization/utils.py index 504d999e858..9948d00ddfa 100644 --- a/megatron/core/inference/quantization/utils.py +++ b/megatron/core/inference/quantization/utils.py @@ -241,9 +241,7 @@ def quantize_params_to_mxfp8( expected_backend=backend, tensor_name=f"persistent MXFP8 parameter {fqn!r}", ) - new_tensor = MXFP8Tensor.from_bf16(bf16_data, backend=backend) - persistent_tensor.data.copy_(new_tensor.data) - persistent_tensor.scale.view(torch.uint8).copy_(new_tensor.scale.view(torch.uint8)) + persistent_tensor.copy_(bf16_data) mcore_tensor = persistent_tensor else: # First call: create new MXFP8Tensor diff --git a/megatron/core/resharding/__init__.py b/megatron/core/resharding/__init__.py index 7fc122c2400..664d925be0e 100644 --- a/megatron/core/resharding/__init__.py +++ b/megatron/core/resharding/__init__.py @@ -1,5 +1,5 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. -from .execution import execute_reshard_plan +from .execution import execute_reshard_plan, refresh_module_caches from .planner import ( build_centralized_reshard_plan, build_local_reshard_plan, @@ -21,6 +21,7 @@ "build_plan_from_rosters", "index_metadata_rosters", "execute_reshard_plan", + "refresh_module_caches", "MXFP8ReshardTransform", "ReshardTransform", "swap_model_weights", diff --git a/megatron/core/resharding/execution.py b/megatron/core/resharding/execution.py index fd28d8cc39f..aec33a94616 100644 --- a/megatron/core/resharding/execution.py +++ b/megatron/core/resharding/execution.py @@ -18,13 +18,23 @@ logger = logging.getLogger(__name__) -def _refresh_module_caches(dst_module: Optional[torch.nn.Module]) -> None: - """Refresh parameter-derived caches in the destination module.""" +def refresh_module_caches( + dst_module: torch.nn.Module | list[torch.nn.Module] | tuple[torch.nn.Module, ...] | None, +) -> None: + """Refresh parameter-derived caches in the destination module(s). + + Lists and tuples let external refit callers, including Megatron Bridge, + pass virtual-pipeline model chunks directly. ``None`` is a no-op for + send-only ranks. The ``dst_module`` parameter name remains stable for + callers that use keyword arguments. + """ if dst_module is None: return - for module in dst_module.modules(): - if isinstance(module, MegatronModule): - module.refresh_cache() + roots = dst_module if isinstance(dst_module, (list, tuple)) else (dst_module,) + for root in roots: + for module in root.modules(): + if isinstance(module, MegatronModule): + module.refresh_cache() @dataclass @@ -103,7 +113,7 @@ def execute_reshard_plan( torch.cuda.synchronize() if service.requires_process_group_barrier: dist.barrier(group=group) - _refresh_module_caches(dst_module) + refresh_module_caches(dst_module) torch.cuda.synchronize() logger.info("Reshard complete") return @@ -253,7 +263,7 @@ def get_sendable(param_name: str, param: torch.nn.Parameter) -> torch.Tensor: dst_param.quantize_(full_bf16) pending_quantized.clear() - _refresh_module_caches(dst_module) + refresh_module_caches(dst_module) # Ensure all writeback and cache-refresh copies are visible to subsequent CUDA # ops (e.g. CUDA graph warmup). The synchronize() above fires *before* the diff --git a/megatron/core/resharding/transforms.py b/megatron/core/resharding/transforms.py index ed5c766d9e8..8f43d3a63cc 100644 --- a/megatron/core/resharding/transforms.py +++ b/megatron/core/resharding/transforms.py @@ -280,9 +280,7 @@ def finalize_recv(self, param_name, dst_slice, recv_buffers): f"1D-scale param {param_name!r}: received {written} elements, " f"expected {buf.data.numel()} (duplicate or missing slices?)" ) - mxfp8 = MXFP8Tensor.from_bf16(accum, backend=buf.backend) - buf.data.copy_(mxfp8.data) - buf.scale.view(torch.uint8).copy_(mxfp8.scale.view(torch.uint8)) + buf.copy_(accum) del self._pending_1d[buf_key] else: self._pending_1d[buf_key][1] = written diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 6b2ccae8466..2b068ff1107 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -1214,6 +1214,7 @@ def _build_concatenated_mxfp8_weights(self): for linear_name, buf_name in [('linear_fc1', '_fc1_weight'), ('linear_fc2', '_fc2_weight')]: linear = getattr(self, linear_name) q_list, s_list = [], [] + source_dtype: torch.dtype | None = None for i in range(self.num_local_experts): w = getattr(linear, f'weight{i}') if isinstance(w, MXFP8Tensor): @@ -1228,6 +1229,13 @@ def _build_concatenated_mxfp8_weights(self): validate_mxfp8_tensor( mxfp8, expected_backend=backend, tensor_name=f"{linear_name}.weight{i}" ) + if mxfp8.dtype is not None: + source_dtype = source_dtype or mxfp8.dtype + if mxfp8.dtype != source_dtype: + raise RuntimeError( + f"Conflicting source dtypes for {linear_name} expert weights: " + f"{source_dtype} and {mxfp8.dtype}." + ) q_list.append(mxfp8.data) s_list.append(mxfp8.scale) @@ -1235,7 +1243,11 @@ def _build_concatenated_mxfp8_weights(self): stacked_scale = torch.stack(s_list, dim=0).contiguous() setattr( - self, buf_name, MXFP8Tensor(data=stacked_data, scale=stacked_scale, backend=backend) + self, + buf_name, + MXFP8Tensor( + data=stacked_data, scale=stacked_scale, dtype=source_dtype, backend=backend + ), ) # Redirect per-expert weight .data to views into the stacked buffer, diff --git a/tests/unit_tests/inference/test_mxfp8_utils.py b/tests/unit_tests/inference/test_mxfp8_utils.py index 595651f3ad1..fa05ecb1ec7 100644 --- a/tests/unit_tests/inference/test_mxfp8_utils.py +++ b/tests/unit_tests/inference/test_mxfp8_utils.py @@ -11,6 +11,8 @@ import pytest import torch +from megatron.core.inference.quantization.mxfp8_tensor import HAVE_FLASHINFER, MXFP8Tensor + pytestmark = [ pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), pytest.mark.internal, @@ -269,6 +271,121 @@ def test_restore_uint8_scale_dtype_after_resharding(self): assert scale.dtype == torch.float8_e8m0fnu torch.testing.assert_close(scale.view(torch.uint8), scale_bytes) + @pytest.mark.parametrize("backend", ["triton", "flashinfer"]) + def test_copy_preserves_storage_and_logical_metadata(self, backend): + if backend == "flashinfer" and not HAVE_FLASHINFER: + pytest.skip("FlashInfer not available") + + source = torch.randn(16, 128, device="cuda", dtype=torch.bfloat16) + tensor = MXFP8Tensor.from_bf16(source, backend=backend) + data_ptr = tensor.data.data_ptr() + scale_ptr = tensor.scale.data_ptr() + + assert tensor.shape == source.shape + assert tensor.dtype == source.dtype + assert tensor.device == source.device + + updated = torch.randn_like(source) + assert tensor.copy_(updated) is tensor + expected = MXFP8Tensor.from_bf16(updated, backend=backend) + + assert tensor.data.data_ptr() == data_ptr + assert tensor.scale.data_ptr() == scale_ptr + assert torch.equal(tensor.data, expected.data) + assert torch.equal(tensor.scale.view(torch.uint8), expected.scale.view(torch.uint8)) + + def test_copy_rejects_shape_mismatch(self): + tensor = MXFP8Tensor.from_bf16( + torch.randn(16, 128, device="cuda", dtype=torch.bfloat16), backend="triton" + ) + with pytest.raises(ValueError, match="shape mismatch"): + tensor.copy_(torch.randn(32, 128, device="cuda", dtype=torch.bfloat16)) + + def test_copy_requires_backend(self): + source = torch.randn(16, 128, device="cuda", dtype=torch.bfloat16) + quantized = MXFP8Tensor.from_bf16(source, backend="triton") + backendless = MXFP8Tensor(data=quantized.data, scale=quantized.scale, dtype=source.dtype) + with pytest.raises(ValueError, match="without a quantization backend"): + backendless.copy_(source) + + def test_copy_preserves_source_dtype_for_legacy_constructor(self): + initial = MXFP8Tensor.from_bf16( + torch.randn(16, 128, device="cuda", dtype=torch.bfloat16), backend="triton" + ) + tensor = MXFP8Tensor(initial.data.clone(), initial.scale.clone(), "triton") + source = torch.randn(16, 128, device="cuda", dtype=torch.float16) + expected = MXFP8Tensor.from_bf16(source, backend="triton") + + assert tensor.dtype is None + tensor.copy_(source) + assert tensor.dtype == source.dtype + assert torch.equal(tensor.data, expected.data) + assert torch.equal(tensor.scale, expected.scale) + + def test_failed_legacy_copy_keeps_dtype_unknown(self): + initial = MXFP8Tensor.from_bf16( + torch.randn(16, 128, device="cuda", dtype=torch.bfloat16), backend="triton" + ) + tensor = MXFP8Tensor( + initial.data.clone(), initial.scale.clone(), "unsupported" # type: ignore[arg-type] + ) + + with pytest.raises(ValueError, match="Unknown MXFP8 quantization backend"): + tensor.copy_(torch.randn(16, 128, device="cuda", dtype=torch.float16)) + assert tensor.dtype is None + + def test_copy_does_not_downcast_update_to_logical_dtype(self): + source = torch.randn(16, 128, device="cuda", dtype=torch.float16) + tensor = MXFP8Tensor.from_bf16(source, backend="triton") + update = torch.linspace( + -100_000, 100_000, steps=16 * 128, device="cuda", dtype=torch.float32 + ).reshape(16, 128) + expected = MXFP8Tensor.from_bf16(update, backend="triton") + + assert tensor.dtype == torch.float16 + tensor.copy_(update) + assert tensor.dtype == torch.float16 + assert torch.equal(tensor.data, expected.data) + assert torch.equal(tensor.scale.view(torch.uint8), expected.scale.view(torch.uint8)) + + def test_copy_normalizes_flashinfer_fp32_input_to_bf16(self, monkeypatch): + tensor = MXFP8Tensor( + data=torch.empty((16, 128), device="cuda", dtype=torch.float8_e4m3fn), + scale=torch.empty(512, device="cuda", dtype=torch.uint8), + backend="flashinfer", + ) + update = torch.randn(16, 128, device="cuda", dtype=torch.float32) + expected_data = torch.zeros_like(tensor.data) + expected_scale = torch.zeros_like(tensor.scale) + quantizer_input = {} + + def fake_from_bf16(cls, value, group_size=32, backend="flashinfer"): + quantizer_input["dtype"] = value.dtype + return cls(data=expected_data, scale=expected_scale, backend=backend, dtype=value.dtype) + + monkeypatch.setattr(MXFP8Tensor, "from_bf16", classmethod(fake_from_bf16)) + + tensor.copy_(update) + + assert quantizer_input["dtype"] == torch.bfloat16 + assert tensor.dtype == torch.bfloat16 + assert torch.equal(tensor.data, expected_data) + assert torch.equal(tensor.scale, expected_scale) + + def test_copy_rejects_stacked_storage(self): + tensor = MXFP8Tensor.from_bf16( + torch.randn(16, 128, device="cuda", dtype=torch.bfloat16), backend="triton" + ) + stacked = MXFP8Tensor( + data=torch.stack([tensor.data, tensor.data]), + scale=torch.stack([tensor.scale, tensor.scale]), + dtype=tensor.dtype, + backend=tensor.backend, + ) + + with pytest.raises(ValueError, match="require 2D destination storage"): + stacked.copy_(torch.randn(2, 16, 128, device="cuda", dtype=torch.bfloat16)) + def test_reject_invalid_scale_dtype(self): from megatron.core.inference.quantization.mxfp8_tensor import ensure_mxfp8_scale_dtype @@ -284,9 +401,9 @@ def test_validate_unswizzled_2d_scale_geometry(self): data = torch.empty((64, 128), dtype=torch.float8_e4m3fn, device="cuda") scale = torch.empty((64, 4), dtype=torch.uint8, device="cuda") - validate_mxfp8_tensor(MXFP8Tensor(data, scale, backend="flashinfer")) + validate_mxfp8_tensor(MXFP8Tensor(data, scale, dtype=torch.bfloat16, backend="flashinfer")) - invalid = MXFP8Tensor(data, scale[:63], backend="flashinfer") + invalid = MXFP8Tensor(data, scale[:63], dtype=torch.bfloat16, backend="flashinfer") with pytest.raises(ValueError, match="2D scale has shape"): validate_mxfp8_tensor(invalid) @@ -300,7 +417,7 @@ def test_missing_backend_error_is_actionable(self): scale = torch.empty(512, dtype=torch.uint8, device="cuda") with pytest.raises(ValueError, match="backend= explicitly"): - validate_mxfp8_tensor(MXFP8Tensor(data, scale)) + validate_mxfp8_tensor(MXFP8Tensor(data, scale, dtype=torch.bfloat16)) @pytest.mark.parametrize("M,K", [(16, 128), (64, 256), (128, 2688)]) def test_from_bf16_triton(self, M, K): @@ -725,6 +842,7 @@ def stack_weights() -> MXFP8Tensor: return MXFP8Tensor( data=torch.stack([weight.data for weight in per_expert]), scale=torch.stack([weight.scale for weight in per_expert]), + dtype=torch.bfloat16, backend="triton", ) diff --git a/tests/unit_tests/resharding/test_cache_refresh.py b/tests/unit_tests/resharding/test_cache_refresh.py index 4d6923abb71..f1231eca937 100644 --- a/tests/unit_tests/resharding/test_cache_refresh.py +++ b/tests/unit_tests/resharding/test_cache_refresh.py @@ -4,7 +4,7 @@ import torch -from megatron.core.resharding.execution import execute_reshard_plan +from megatron.core.resharding.execution import execute_reshard_plan, refresh_module_caches from megatron.core.resharding.utils import ReshardPlan from megatron.core.transformer.module import MegatronModule @@ -24,6 +24,15 @@ def test_default_refresh_cache_is_noop(): assert module.refresh_cache() is None +def test_refresh_module_caches_accepts_model_lists(): + events: list[str] = [] + + refresh_module_caches(None) + refresh_module_caches([torch.nn.Sequential(_CacheAwareModule(events)), torch.nn.Sequential()]) + + assert events == ["refresh"] + + def test_execute_reshard_plan_refreshes_caches_before_final_sync(monkeypatch): events: list[str] = [] target_core = torch.nn.Sequential(_CacheAwareModule(events)) diff --git a/tests/unit_tests/resharding/test_mxfp8_refit.py b/tests/unit_tests/resharding/test_mxfp8_refit.py index 3506dbc1fdd..a97511ff08c 100644 --- a/tests/unit_tests/resharding/test_mxfp8_refit.py +++ b/tests/unit_tests/resharding/test_mxfp8_refit.py @@ -188,6 +188,7 @@ def test_mixed_persistent_buffer_backends_are_rejected(self): flashinfer_buffer = MXFP8Tensor( data=triton_buffer.data.clone(), scale=triton_buffer.scale.view(torch.uint8).clone(), + dtype=torch.bfloat16, backend="flashinfer", ) @@ -218,6 +219,8 @@ class Namespace: tensor = MXFP8Tensor.from_bf16( torch.randn(M, K, dtype=torch.bfloat16, device="cuda"), backend="triton" ) + if expert_idx == 0: + tensor.dtype = None # Simulate a legacy direct constructor. setattr(linear, f"weight{expert_idx}", tensor) buffers[f"{linear_name}.weight{expert_idx}"] = tensor @@ -228,6 +231,8 @@ class Namespace: assert not torch.is_inference(grouped_mlp._fc1_weight.scale) assert not torch.is_inference(grouped_mlp._fc2_weight.data) assert not torch.is_inference(grouped_mlp._fc2_weight.scale) + assert grouped_mlp._fc1_weight.dtype == torch.bfloat16 + assert grouped_mlp._fc2_weight.dtype == torch.bfloat16 transform = MXFP8ReshardTransform( convertible_params=set(buffers), persistent_buffers=buffers