From 1304ac5743510912e64461fc39657ebcb86778b4 Mon Sep 17 00:00:00 2001 From: Ivy Zhou Date: Fri, 21 Aug 2026 18:20:20 -0700 Subject: [PATCH] [float8] Use stateless grouped MM for MoE training Summary: Keep Float8 grouped-expert parameters as ordinary tensors and route the existing grouped-matmul seam through TorchAO stateless dynamic quantization. This avoids serializing TorchAO wrapper/config objects while preserving the FP8 computation path. Test Plan: - pytest -q tests/unit_tests/test_quantization.py with current TorchAO and 0.18.0 - pre-commit run --all-files - H100 eager and fullgraph-compiled forward/backward parity - 2-GPU FSDP2 + DCP save/load forward/backward parity --- tests/unit_tests/test_quantization.py | 47 ++++++++++++++++++++ torchtitan/components/quantization/float8.py | 14 +++--- 2 files changed, 56 insertions(+), 5 deletions(-) diff --git a/tests/unit_tests/test_quantization.py b/tests/unit_tests/test_quantization.py index b4a17d023a..e2a4d0124d 100644 --- a/tests/unit_tests/test_quantization.py +++ b/tests/unit_tests/test_quantization.py @@ -6,6 +6,7 @@ import pytest import spmd_types as spmd import torch +import torch.distributed.checkpoint as dcp from torchtitan.components.data import ( FirstFitPackingConfig, @@ -365,3 +366,49 @@ def test_quantized_grouped_experts(): assert issubclass(float8_cls, GptOssGroupedExperts) assert hasattr(mxfp8_cls.Config, "swiglu_limit") assert hasattr(float8_cls.Config, "swiglu_limit") + + +@pytest.mark.parametrize("parent_cls", [GroupedExperts, GptOssGroupedExperts]) +def test_float8_grouped_experts_checkpoint_state_uses_plain_tensors(parent_cls): + pytest.importorskip("torchao") + stock = parent_cls.Config(dim=16, hidden_dim=32, num_experts=2).build() + float8_cls = _get_float8_grouped_experts_cls(parent_cls) + module = float8_cls.Config(dim=16, hidden_dim=32, num_experts=2).build() + + assert all(type(param) is torch.nn.Parameter for param in module.parameters()) + stock_state = stock.state_dict() + float8_state = module.state_dict() + assert float8_state.keys() == stock_state.keys() + for key, value in float8_state.items(): + assert type(value) is torch.Tensor + assert value.shape == stock_state[key].shape + assert value.dtype == stock_state[key].dtype + + +@pytest.mark.filterwarnings("ignore:torch.distributed is disabled") +def test_float8_grouped_experts_dcp_round_trip_needs_no_safe_globals(tmp_path): + pytest.importorskip("torchao") + float8_cls = _get_float8_grouped_experts_cls(GroupedExperts) + config = float8_cls.Config(dim=16, hidden_dim=32, num_experts=2) + source = config.build() + target = config.build() + + with torch.no_grad(): + for value, parameter in enumerate(source.parameters(), start=1): + parameter.fill_(value) + for parameter in target.parameters(): + parameter.zero_() + + saved_safe_globals = torch.serialization.get_safe_globals() + try: + torch.serialization.clear_safe_globals() + dcp.save(source.state_dict(), checkpoint_id=tmp_path, no_dist=True) + dcp.load(target.state_dict(), checkpoint_id=tmp_path, no_dist=True) + finally: + torch.serialization.clear_safe_globals() + torch.serialization.add_safe_globals(saved_safe_globals) + + for source_parameter, target_parameter in zip( + source.parameters(), target.parameters(), strict=True + ): + torch.testing.assert_close(target_parameter, source_parameter) diff --git a/torchtitan/components/quantization/float8.py b/torchtitan/components/quantization/float8.py index 791a4406d6..99770ed0c9 100644 --- a/torchtitan/components/quantization/float8.py +++ b/torchtitan/components/quantization/float8.py @@ -195,12 +195,16 @@ class Config(parent_config_cls): # type: ignore[misc] def __init__(self, config: Config): super().__init__(config) from torchao.prototype.moe_training.config import Float8TrainingOpConfig - from torchao.quantization.quant_api import quantize_ - quantize_( - self, - config=Float8TrainingOpConfig(), - filter_fn=lambda mod, _fqn: isinstance(mod, GroupedExperts), + self._float8_op_config = Float8TrainingOpConfig() + + def _grouped_mm(self, *, A, B_t, offs): + from torchao.prototype.moe_training.utils import ( + _quantize_then_scaled_grouped_mm, + ) + + return _quantize_then_scaled_grouped_mm( + A, B_t, config=self._float8_op_config, offs=offs ) Float8GroupedExperts.__name__ = f"Float8{parent_cls.__name__}"