float8 moe training conversion API prototype

stack-info: PR: #2275, branch: danielvegamyhre/stack/1

migrate to quantize and add test

work on moe training test

Author: danielvegamyhre

test/prototype/scaled_grouped_mm/test_moe_training_conversion.py

@@ -0,0 +1,84 @@
+import pytest
+import torch
+from torch import nn
+
+from torchao.float8.float8_utils import compute_error
+from torchao.prototype.scaled_grouped_mm.conversion_utils import MoETrainingConfig
+from torchao.quantization.quant_api import quantize_
+
+try:
+    from torchtitan.experiments.llama4.model.args import TransformerModelArgs
+    from torchtitan.experiments.llama4.model.moe import MoE
+except ImportError:
+    import warnings
+
+    warnings.warn("torchtitan not installed, skipping MoE tests.")
+    pytest.skip(allow_module_level=True)
+
+
+@pytest.mark.parametrize(
+    "target_fqns",
+    ["experts"],
+)
+def test_moe_float8_training(target_fqns: list[str]):
+    model_args = TransformerModelArgs(moe_enabled=True, num_experts=2)
+    init_std = 0.02
+    device = torch.device("cuda")
+
+    # reference bf16 MoE
+    ref_model = MoE(model_args).to(torch.bfloat16).cuda()
+    torch.manual_seed(42)
+    with torch.no_grad():
+        ref_model.init_weights(init_std, device)
+
+    # target MoE for testing conversion
+    model = MoE(model_args).to(torch.bfloat16).cuda()
+    torch.manual_seed(42)
+    with torch.no_grad():
+        model.init_weights(init_std, device)
+
+    # assert starting params are identical for both models
+    for param1, param2 in zip(model.parameters(), ref_model.parameters()):
+        assert torch.equal(param1, param2)
+
+    # convert MoE to float8 training
+    def moe_module_filter_fn(mod: nn.Module, cur_fqn: str) -> bool:
+        for target_fqn in target_fqns:
+            if target_fqn in cur_fqn:
+                return True
+        return False
+
+    config = MoETrainingConfig()
+    quantize_(model, config=config, filter_fn=moe_module_filter_fn)
+
+    # inputs
+    batch, seq, dim = 1, 8192, 4096
+    x = torch.randn(batch, seq, dim, dtype=torch.bfloat16, requires_grad=True).cuda()
+    ref_x = x.clone()
+
+    # forward pass
+    out = model(x)
+    ref_out = ref_model(ref_x)
+
+    # validate SQNR between outputs is acceptable.
+    # a single fp8 gemm uses SQNR >= 25.0 for testing, so for a full MoE layer
+    # we'll use a slightly lower threshold.
+    out_sqnr = compute_error(out, ref_out)
+    assert out_sqnr.item() >= 23.0, f"SQNR must be >= 23.0, got {out_sqnr.item()}."
+
+    # backward pass
+    ref_out.sum().backward()
+    out.sum().backward()
+
+    # validate input gradients
+    input_grad_sqnr = compute_error(x.grad, ref_x.grad)
+    assert input_grad_sqnr.item() >= 23.0, (
+        f"SQNR must be >= 23.0, got {input_grad_sqnr.item()}."
+    )
+
+    # validate param gradients
+    for param1, param2 in zip(model.parameters(), ref_model.parameters()):
+        param_grad_sqnr = compute_error(param1.grad, param2.grad)
+        assert param_grad_sqnr.item() >= 23.0, (
+            f"SQNR must be >= 23.0, got {param_grad_sqnr.item()}."
+        )

test/prototype/scaled_grouped_mm/test_scaled_grouped_mm.py

@@ -23,7 +23,7 @@
     Float8LinearConfig,
     Float8LinearRecipeName,
 )
-from torchao.float8.float8_linear import matmul_with_hp_or_float8_args
+from torchao.float8.float8_linear import _matmul_with_hp_or_float8_args
 from torchao.float8.float8_tensor import LinearMMConfig
 from torchao.float8.float8_utils import tensor_to_scale, to_fp8_saturated
 from torchao.prototype.scaled_grouped_mm.scaled_grouped_mm import (
@@ -183,7 +183,7 @@ def compute_reference_forward(
 
     # Validate each actual result group from the _scaled_grouped_mm is equal to:
     # 1. A manual _scaled_mm for the group.
-    # 2. A matmul_with_hp_or_float8_args for the group (which is differentiable, and thus used to validate gradients).
+    # 2. A _matmul_with_hp_or_float8_args for the group (which is differentiable, and thus used to validate gradients).
     outputs = []
     list1 = list(zip(A_list_fp8, B_t_fp8, A_scale_list, B_t_scales, result_list))
     list2 = list(zip(A_list, B_t, result_list))
@@ -199,7 +199,7 @@ def compute_reference_forward(
             use_fast_accum=float8_config.gemm_config_output.use_fast_accum,
         )
         a2, b2, result2 = list2[i]
-        ref_group_result2 = matmul_with_hp_or_float8_args.apply(
+        ref_group_result2 = _matmul_with_hp_or_float8_args.apply(
             a2,
             b2,
             LinearMMConfig(),

torchao/prototype/scaled_grouped_mm/conversion_utils.py

@@ -0,0 +1,93 @@
+from typing import Callable, Optional
+
+from torch import nn
+
+from torchao.core.config import AOBaseConfig
+from torchao.quantization.transform_module import (
+    register_quantize_module_handler,
+)
+from torchao.prototype.scaled_grouped_mm.tensor import ScaledGroupedMMTensor
+
+
+class MoETrainingConfig(AOBaseConfig):
+    pass
+
+
+@register_quantize_module_handler(MoETrainingConfig)
+def _moe_training_transform(
+    module: nn.Module,
+    config: MoETrainingConfig,
+) -> nn.Module:
+    """
+    Swaps `torch.nn.Parameter` data tensor with a ScaledGroupedMMTensor.
+
+    Args:
+        module: Module to modify.
+        config: MoETrainingConfig which defines how to perform the MoE training transform.
+
+    Returns:
+     nn.Module: The modified module with swapped parameters.
+    """
+    out = swap_params(module)
+    return out
+
+def swap_params(
+    module: nn.Module,
+    *,
+    module_filter_fn: Optional[Callable[[nn.Module, str], bool]] = None,
+) -> nn.Module:
+    """
+    Recurses through the nn.Module, recursively swapping the data tensor of
+    each nn.Parameter with a ScaledGroupedMMTensor. Only applies if the module
+    passed the module_filter_fn, if specified.
+
+    Args:
+        module: Module to modify.
+        module_filter_fn: If specified, only the `torch.nn.Parameter` subclasses that
+            that pass the filter function will be swapped. The inputs to the
+            filter function are the module instance, and the FQN.
+
+    Returns:
+     nn.Module: The modified module with swapped linear layers.
+    """
+    if isinstance(module, nn.Parameter) and (
+        module_filter_fn is None or module_filter_fn(module, "")
+    ):
+        if len(list(module.children())) > 0:
+            raise AssertionError(
+                f"Does not support a root nn.Parameter with children: {module}"
+            )
+        if not isinstance(module.data, ScaledGroupedMMTensor):
+            new_data = ScaledGroupedMMTensor(module.data)
+            return nn.Parameter(new_data, requires_grad=module.requires_grad)
+        return module
+
+    root_module = module
+
+    def post_order_traversal(
+        module: nn.Module,
+        cur_fqn: Optional[str] = None,
+        parent_module: Optional[nn.Module] = None,
+    ):
+        if cur_fqn is None:
+            cur_fqn = ""
+
+        for child_module_name, child_module in module.named_children():
+            if cur_fqn == "":
+                new_fqn = child_module_name
+            else:
+                new_fqn = f"{cur_fqn}.{child_module_name}"
+
+            post_order_traversal(child_module, new_fqn, module)
+
+        if module_filter_fn is None or module_filter_fn(module, cur_fqn):
+            for param_name, param in module.named_parameters(recurse=False):
+                if not isinstance(param.data, ScaledGroupedMMTensor):
+                    new_param = nn.Parameter(
+                        ScaledGroupedMMTensor(param), requires_grad=param.requires_grad
+                    )
+                    setattr(module, param_name, new_param)
+                    print(f"Swapped {cur_fqn}.{param_name} to ScaledGroupedMMTensor")
+
+    post_order_traversal(root_module)
+    return root_module

torchao/prototype/scaled_grouped_mm/scaled_grouped_mm.py

@@ -83,7 +83,10 @@ def forward(
         assert not _is_column_major(A), "A must be row-major"
 
         # Due to hardware requirements, the right operand in a scaled grouped GEMM must be column-major.
-        assert _is_column_major(B_t), "B must be column-major"
+        if not _is_column_major(B_t):
+            # FSDP will complain if B_t (weights) is not contiguous, we can't require B_t to be column-major.
+            # TODO: figure out better solution than transposing for each forward pass.
+            B_t = B_t.transpose(-2, -1).contiguous().transpose(-2, -1)
 
         # Convert high precision input tensor to float8, row-major for left operand of grouped GEMM.
         # A shape: (M, K)

torchao/prototype/scaled_grouped_mm/tensor.py

@@ -0,0 +1,35 @@
+import torch
+
+from torchao.prototype.scaled_grouped_mm import _scaled_grouped_mm
+
+
+class ScaledGroupedMMTensor(torch.Tensor):
+    """
+    ScaledGroupedMMTensor is a simple tensor subclass that wraps a regular tensor
+    and overrides the torch._grouped_mm op by dispatching to the
+    differentiable _scaled_grouped_mm autograd function.
+    """
+
+    grouped_mm_func_name = "_grouped_mm"
+    offs_arg_name = "offs"
+
+    def __init__(self, data: torch.Tensor):
+        self._data = data
+
+    @classmethod
+    def __torch_function__(cls, func, types, args, kwargs={}):
+        if func.__name__ == cls.grouped_mm_func_name:
+            # Use torchao scaled grouped mm with dynamic quant for
+            # "2d x 3d with offsets" case (used for routed experts).
+            # Otherwise, fall back to regular grouped mm.
+            #
+            # TODO: support "3d x 3d without offsets" case, which is
+            # used for shared experts. This is basically the grouped_mm
+            # kernel handling a bmm.
+            A, B = args[0], args[1]
+            A_is_2d = A.dim() == 2
+            B_is_3d = B.dim() == 3
+            has_offs = kwargs.get(cls.offs_arg_name) is not None
+            if A_is_2d and B_is_3d and has_offs:
+                return _scaled_grouped_mm(*args, **kwargs)
+        return super().__torch_function__(func, types, args, kwargs)