[mxfp8 moe training] mxfp8 a2a working e2e in torchtitan llama4 training; improve tests + bench scripts (#3088)

[mxfp8 moe training] fix CUDA IMA and improve bench + test scripts

Author: danielvegamyhre

benchmarks/prototype/moe_training/mxfp8/bench_all_to_all_v.py

@@ -10,6 +10,7 @@
 # torchrun --nproc-per-node=8 --local-ranks-filter=0 benchmarks/prototype/moe_training/mxfp8/bench_all_to_all_v.py
 #
 #######################################################################
+import argparse
 import os
 import time
 from dataclasses import dataclass
@@ -18,11 +19,14 @@
 import torch
 from tabulate import tabulate
 from torch import distributed as dist
+from torch.distributed import DeviceMesh, init_device_mesh
 from torch.distributed._functional_collectives import (
+    all_to_all_single,
     all_to_all_single_autograd,
 )
 from tqdm import tqdm
 
+from benchmarks.utils import profile_fn
 from torchao.prototype.moe_training.kernels.mxfp8.comms import (
     mxfp8_on_device_all_to_all_v,
 )
@@ -37,8 +41,8 @@ class ExperimentConfig:
 
 @dataclass(frozen=True)
 class ExperimentResult:
-    bf16_us: float
-    mxfp8_us: float
+    bf16_ms: float
+    mxfp8_ms: float
 
 
 @dataclass(frozen=True)
@@ -50,7 +54,7 @@ class Experiment:
 def get_configs() -> List[ExperimentConfig]:
     # (batch_size, seq_len, dim)
     input_shapes = [
-        (8, 8192, 5120),
+        (16, 8192, 5120),
     ]
     configs = []
     for shape in input_shapes:
@@ -62,7 +66,111 @@ def get_configs() -> List[ExperimentConfig]:
     return configs
 
 
-def run_experiment(config: ExperimentConfig) -> ExperimentResult:
+# Copy/paste a2a impls added in https://github.com/pytorch/torchtitan/pull/1765
+def default_a2a_dispatch(
+    routed_input: torch.Tensor,
+    num_tokens_per_expert: torch.Tensor,
+    device_mesh: DeviceMesh,
+):
+    """
+    Default implementation of all-to-all dispatch. Incurs device-to-host sync.
+
+    Returns:
+        routed_input: the local tokens after all-to-all dispatch
+        input_splits: the input splits for all-to-all dispatch
+        output_splits: the output splits for all-to-all dispatch
+        num_tokens_per_expert_group: the number of tokens per EP rank after all-to-all dispatch
+    """
+    ep_degree = device_mesh.size(0)
+    # generate the input splits and output splits for all-to-all
+    with torch.no_grad():
+        num_tokens_per_expert_group = all_to_all_single(
+            num_tokens_per_expert,
+            None,
+            None,
+            group=device_mesh.get_group(),
+        )
+        # Need to wait explicitly because it is used by a triton kernel later
+        # which doesn't realize that AsyncCollectiveTensor needs unwrapping
+        num_tokens_per_expert_group = torch.ops._c10d_functional.wait_tensor(
+            num_tokens_per_expert_group
+        )
+        input_splits = (
+            num_tokens_per_expert.view(ep_degree, -1)
+            .sum(dim=1)
+            .to(torch.device("cpu"), non_blocking=True)
+        )
+        # NOTE: this would incur a device-to-host sync
+        output_splits = (
+            num_tokens_per_expert_group.view(ep_degree, -1)
+            .sum(dim=1)
+            .to(torch.device("cpu"), non_blocking=False)
+        )
+        input_splits_list = input_splits.tolist()
+        output_splits_list = output_splits.tolist()
+
+    # perform all-to-all
+    routed_input = all_to_all_single_autograd(
+        routed_input,
+        output_splits_list,
+        input_splits_list,
+        device_mesh.get_group(),
+    )
+    routed_input = torch.ops._c10d_functional.wait_tensor(routed_input)
+    return (
+        routed_input,
+        input_splits_list,
+        output_splits_list,
+        num_tokens_per_expert_group,
+    )
+
+
+def mxfp8_a2a_dispatch(
+    routed_input: torch.Tensor,
+    num_tokens_per_expert: torch.Tensor,
+    device_mesh: DeviceMesh,
+    max_tokens_per_ep_rank: int,
+):
+    """
+    Perform on-device all-to-all dispatch with dynamically quantized mxfp8 inputs to save network bandwidth
+    and avoid device-to-host sync.
+
+    Returns:
+        routed_input: the local tokens after all-to-all dispatch
+        input_splits: the input splits for all-to-all dispatch
+        output_splits: the output splits for all-to-all dispatch
+    """
+
+    ep_degree = device_mesh.size(0)
+    num_tokens_per_expert_group = all_to_all_single(
+        num_tokens_per_expert,
+        None,
+        None,
+        group=device_mesh.get_group(),
+    )
+    input_splits_per_ep_rank = num_tokens_per_expert.view(ep_degree, -1).sum(dim=1)
+    num_tokens_per_expert_group = torch.ops._c10d_functional.wait_tensor(
+        num_tokens_per_expert_group
+    )
+    routed_input, output_splits_per_ep_rank = mxfp8_on_device_all_to_all_v(
+        routed_input,
+        input_splits_per_ep_rank,
+        max_tokens_per_ep_rank,
+        device_mesh.get_group().group_name,
+    )
+    tokens_on_rank_after_a2a = output_splits_per_ep_rank.sum()
+    routed_input_no_padding = routed_input[:tokens_on_rank_after_a2a]
+    return (
+        routed_input_no_padding,
+        input_splits_per_ep_rank,
+        output_splits_per_ep_rank,
+        num_tokens_per_expert_group,
+    )
+
+
+def run_experiment(
+    config: ExperimentConfig, args: argparse.Namespace
+) -> ExperimentResult:
     batch_size, seq_len, dim = config.input_shape
     x = torch.randn(
         (batch_size * seq_len, dim),
@@ -71,99 +179,70 @@ def run_experiment(config: ExperimentConfig) -> ExperimentResult:
     )
     ref_x = x.detach().clone()
 
+    # Set up device mesh
+    mesh = init_device_mesh("cuda", (dist.get_world_size(),))
+
     # Max output tokens per rank is worst case where one rank receives all tokens
     input_tokens_per_rank = batch_size * seq_len
     max_output_tokens_per_rank = input_tokens_per_rank * dist.get_world_size()
 
-    def using_bf16(
-        input_tensor: torch.Tensor, input_splits: torch.Tensor
-    ) -> torch.Tensor:
-        # Calculate output splits from input splits
-        output_splits = torch.empty_like(input_splits)
-        dist.all_to_all_single(output_splits, input_splits)
-
-        # Perform all-to-all
-        out = all_to_all_single_autograd(
-            input_tensor,
-            output_splits.tolist(),
-            input_splits.tolist(),
-            dist.group.WORLD,
-        )
-        out = torch.ops._c10d_functional.wait_tensor(out)
-        return out
-
-    def using_mxfp8(
-        input_tensor: torch.Tensor, input_splits: torch.Tensor
-    ) -> torch.Tensor:
-        output, output_splits = mxfp8_on_device_all_to_all_v(
-            input_tensor,
-            input_splits,
-            max_output_tokens_per_rank,
-            dist.group.WORLD.group_name,
-        )
-        output = torch.ops._c10d_functional.wait_tensor(output)
-        output_splits = torch.ops._c10d_functional.wait_tensor(output_splits)
-        return output
-
     def warmup(func_no_args):
         for _ in range(2):
             func_no_args()
 
-    num_splits = dist.get_world_size()
+    num_experts_per_rank = 2
+    num_splits = dist.get_world_size() * num_experts_per_rank
     input_splits = generate_split_sizes(
         num_splits, input_tokens_per_rank, device=device
     )
 
-    print(
-        "Benchmarking using bf16",
-        "batch_size",
-        batch_size,
-        "seq_len",
-        seq_len,
-        "dim",
-        dim,
-        "input_tokens_per_rank",
-        input_tokens_per_rank,
-        "max_output_tokens_per_rank",
-        max_output_tokens_per_rank,
-    )
-    warmup(lambda: using_bf16(ref_x, input_splits))
-    start_ns = time.perf_counter()
-    using_bf16(ref_x, input_splits)
-    end_ns = time.perf_counter()
-    bf16_us = (end_ns - start_ns) * 1e6
-
-    print(
-        "Benchmarking using_mxfp8",
-        "batch_size",
-        batch_size,
-        "seq_len",
-        seq_len,
-        "dim",
-        dim,
-        "input_tokens_per_rank",
-        input_tokens_per_rank,
-        "max_output_tokens_per_rank",
-        max_output_tokens_per_rank,
+    # Bench default a2a
+    warmup(lambda: default_a2a_dispatch(ref_x, input_splits, mesh))
+    start_sec = time.perf_counter()
+    default_a2a_dispatch(ref_x, input_splits, mesh)
+    end_sec = time.perf_counter()
+    bf16_ms = (end_sec - start_sec) * 1e3
+    if args.profile:
+        profile_fn(
+            default_a2a_dispatch,
+            ref_x,
+            input_splits,
+            mesh,
+            distributed=True,
+            profile_name="all_to_all_single_autograd",
+        )
+
+    # Bench mxfp8 a2a
+    warmup(
+        lambda: mxfp8_a2a_dispatch(x, input_splits, mesh, max_output_tokens_per_rank)
     )
-    warmup(lambda: using_mxfp8(x, input_splits))
-    start_ns = time.perf_counter()
-    using_mxfp8(x, input_splits)
-    end_ns = time.perf_counter()
-    mxfp8_us = (end_ns - start_ns) * 1e6
+    start_sec = time.perf_counter()
+    mxfp8_a2a_dispatch(x, input_splits, mesh, max_output_tokens_per_rank)
+    end_sec = time.perf_counter()
+    mxfp8_ms = (end_sec - start_sec) * 1e3
+    if args.profile:
+        profile_fn(
+            mxfp8_a2a_dispatch,
+            x,
+            input_splits,
+            mesh,
+            max_output_tokens_per_rank,
+            distributed=True,
+            profile_name="mxfp8_all_to_all_v",
+        )
 
     return ExperimentResult(
-        bf16_us=bf16_us,
-        mxfp8_us=mxfp8_us,
+        bf16_ms=bf16_ms,
+        mxfp8_ms=mxfp8_ms,
     )
 
 
 def print_results(experiments: List[Experiment]):
     headers = [
         "input_shape",
         "num_splits",
-        "bf16_us",
-        "mxfp8_us",
+        "bf16_ms",
+        "mxfp8_ms",
     ]
     rows = []
     num_splits = dist.get_world_size()
@@ -172,8 +251,8 @@ def print_results(experiments: List[Experiment]):
             [
                 str(experiment.config.input_shape),
                 num_splits,
-                experiment.result.bf16_us,
-                experiment.result.mxfp8_us,
+                experiment.result.bf16_ms,
+                experiment.result.mxfp8_ms,
             ]
         )
     print(tabulate(rows, headers=headers))
@@ -209,7 +288,7 @@ def generate_split_sizes(K: int, N: int, device: str = "cuda") -> torch.Tensor:
     return result.to(dtype=torch.int64)
 
 
-def main():
+def main(args: argparse.Namespace):
     torch.random.manual_seed(123)
 
     # Set up process group
@@ -219,7 +298,7 @@ def main():
     configs = get_configs()
     results = []
     for config in tqdm(configs):
-        result = run_experiment(config)
+        result = run_experiment(config, args)
         results.append(Experiment(config=config, result=result))
 
     # Use Tabulate to print results
@@ -237,4 +316,7 @@ def setup_distributed():
 
 
 if __name__ == "__main__":
-    main()
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--profile", action="store_true")
+    args = parser.parse_args()
+    main(args)

benchmarks/utils.py

@@ -72,7 +72,7 @@ def profile_fwd_bwd(
     print(f"Saved: {profile_name}.json")
 
 
-def profile_fn(fn, *args, profile_name="profile", **kwargs):
+def profile_fn(fn, *args, profile_name="profile", distributed=False, **kwargs):
     wait, warmup, active = 1, 1, 1
     total_steps = wait + warmup + active
     with torch.profiler.profile(
@@ -89,9 +89,11 @@ def profile_fn(fn, *args, profile_name="profile", **kwargs):
             _ = fn(*args, **kwargs)
             prof.step()
 
-    # Save profiler results
-    prof.export_chrome_trace(f"{profile_name}.json")
-    print(f"Saved: {profile_name}.json")
+    if distributed:
+        if torch.distributed.get_rank() == 0:
+            # Save profiler results
+            prof.export_chrome_trace(f"{profile_name}.json")
+            print(f"Saved: {profile_name}.json")
 
 
 def benchmark_cuda_function_in_microseconds(f, *args, **kwargs):

test/prototype/moe_training/mxfp8/test_mxfp8_a2a.py

@@ -29,7 +29,7 @@
 
 
 @instantiate_parametrized_tests
-class TritonAllReduceTest(MultiProcessTestCase):
+class MXFP8AllToAllVTest(MultiProcessTestCase):
     def setUp(self) -> None:
         super().setUp()
         self._spawn_processes()