[WIP] separate matmul and all reduce

Author: linamy85

Ironwood/configs/training/local_separation_test.yaml

@@ -0,0 +1,34 @@
+benchmarks:
+  - benchmark_name: gemm_only
+    trace_dir: "../microbenchmarks/gemm_all_reduce"
+    csv_path: "../microbenchmarks/gemm_all_reduce"
+    xlml_metrics_dir: "../microbenchmarks/gemm_all_reduce"
+    xla_dump_dir: "../microbenchmarks/gemm_all_reduce/hlo_graphs"
+    num_runs: 10
+    benchmark_params:
+      - m: 4096
+        k: 4096
+        n: 4096
+        dtype: "bfloat16"
+  - benchmark_name: all_reduce_only
+    trace_dir: "../microbenchmarks/gemm_all_reduce"
+    csv_path: "../microbenchmarks/gemm_all_reduce"
+    xlml_metrics_dir: "../microbenchmarks/gemm_all_reduce"
+    xla_dump_dir: "../microbenchmarks/gemm_all_reduce/hlo_graphs"
+    num_runs: 10
+    benchmark_params:
+      - m: 4096
+        k: 4096  # Passed to maintain signature, though not used for shape of C
+        n: 4096
+        dtype: "bfloat16"
+  - benchmark_name: gemm_all_reduce
+    trace_dir: "../microbenchmarks/gemm_all_reduce"
+    csv_path: "../microbenchmarks/gemm_all_reduce"
+    xlml_metrics_dir: "../microbenchmarks/gemm_all_reduce"
+    xla_dump_dir: "../microbenchmarks/gemm_all_reduce/hlo_graphs"
+    num_runs: 10
+    benchmark_params:
+      - m: 4096
+        k: 4096
+        n: 4096
+        dtype: "bfloat16"

Ironwood/src/benchmark_gemm_all_reduce.py

@@ -16,7 +16,8 @@
     MetricsStatistics,
     get_metrics_helper,
     str_to_dtype,
-    get_peak_flops_multiplier
+    get_peak_flops_multiplier,
+    unified_bytes_metrics,
 )
 from common import MARKER
 import jax
@@ -178,18 +179,6 @@ def gemm_all_reduce_calculate_metrics(
     peak_flops_multiplier = get_peak_flops_multiplier(dtype_str)
     peak_flops = PEAK_FLOPS_PER_DEVICE * peak_flops_multiplier
 
-    # Unified FLOPs metrics
-    metadata, metrics = unified_flops_metrics(
-        m,
-        n,
-        k,
-        time_ms_list,
-        total_flops_per_device,
-        total_flops_all_devices,
-        peak_flops,
-        dtype=dtype_str,
-    )
-
     # Calculate Bandwidth for AllReduce
     # AllReduce moves Matrix C: M x N
     matrix_c_size_bytes = m * n * dtype.dtype.itemsize
@@ -201,3 +190,227 @@ def gemm_all_reduce_calculate_metrics(
     )
 
     return metadata, metrics
+
+
+def gemm_only(
+    m: int,
+    k: int,
+    n: int,
+    dtype: jnp.dtype = jnp.bfloat16,
+    num_runs: int = 1,
+    trace_dir: str = None,
+) -> Dict[str, Any]:
+    """Benchmarks only the Matmul part of gemm_all_reduce.
+    
+    A: [M, K]
+    B: [K, N]
+    C = A @ B: [M, N]
+    """
+
+    dtype_str = dtype.dtype.name
+    print(f"Running gemm_only benchmark with m={m}, k={k}, n={n}, dtype={dtype_str}, runs={num_runs}")
+
+    def f(x, y):
+        with jax.named_scope(MARKER):
+            # Matmul
+            acc = jax.numpy.einsum(
+                "ij,jk->ik", x, y, preferred_element_type=jnp.float32
+            )
+            c = acc.astype(dtype)
+            return c
+
+    mesh = create_mesh(SHARDING_STRATEGY)
+    lhs_sharding = get_lhs_named_shading(mesh, SHARDING_STRATEGY)
+    rhs_sharding = get_rhs_named_shading(mesh, SHARDING_STRATEGY)
+    out_sharding = get_out_sharding(SHARDING_STRATEGY)
+
+    jit_sharded_f = jax.jit(
+        shard_map(
+            f,
+            mesh,
+            in_specs=(
+                lhs_sharding.spec,
+                rhs_sharding.spec,
+            ),
+            out_specs=out_sharding,
+            check_rep=False,
+        )
+    )
+
+    lhs_shape = (m, k)
+    rhs_shape = (k, n)
+
+    lhs_dtype = dtype
+    rhs_dtype = dtype
+
+    key = jax.random.key(SEED)
+
+    def data_generator():
+        """Creates new random data on host and puts it on device."""
+        nonlocal key
+        key, key_lhs, key_rhs = jax.random.split(key, 3)
+
+        # Create random data on host
+        lhs_host = jax.random.normal(key_lhs, lhs_shape).astype(lhs_dtype)
+        rhs_host = jax.random.normal(key_rhs, rhs_shape).astype(rhs_dtype)
+
+        # Put on device (HBM)
+        lhs_device = jax.device_put(lhs_host, lhs_sharding)
+        rhs_device = jax.device_put(rhs_host, rhs_sharding)
+
+        return (lhs_device, rhs_device)
+
+    time_ms_list = iteration_timeit(
+        jit_sharded_f,
+        data_generator,
+        matrix_dim=f"{dtype_str}_{m}x{n}x{k}",
+        tries=num_runs,
+        task=f"gemm_only_{dtype_str}",
+        trace_dir=trace_dir,
+    )
+    return {
+        "time_ms_list": time_ms_list,
+    }
+
+
+def gemm_only_calculate_metrics(
+    m: int,
+    k: int,
+    n: int,
+    dtype: jnp.dtype,
+    time_ms_list: list[float],
+) -> Dict[str, Any]:
+    # Calculate FLOPs (Matmul)
+    total_flops = 2 * m * k * n
+    
+    total_flops_per_device, total_flops_all_devices = handle_based_on_sharding(
+        total_flops, SHARDING_STRATEGY
+    )
+    
+    dtype_str = dtype.dtype.name
+    peak_flops_multiplier = get_peak_flops_multiplier(dtype_str)
+    peak_flops = PEAK_FLOPS_PER_DEVICE * peak_flops_multiplier
+    
+    metadata, metrics = unified_flops_metrics(
+        m, n, k, time_ms_list, total_flops_per_device, total_flops_all_devices, peak_flops, dtype=dtype_str,
+    )
+    
+    return metadata, metrics
+
+
+def all_reduce_only(
+    m: int,
+    k: int,
+    n: int,
+    dtype: jnp.dtype = jnp.bfloat16,
+    num_runs: int = 1,
+    trace_dir: str = None,
+) -> Dict[str, Any]:
+    """Benchmarks only the AllReduce part of gemm_all_reduce independently.
+    
+    Input: C [M, N]
+    Output = AllReduce(C)
+    """
+
+    dtype_str = dtype.dtype.name
+    print(f"Running all_reduce_only benchmark with m={m}, k={k}, n={n}, dtype={dtype_str}, runs={num_runs}")
+
+    def f(c):
+        with jax.named_scope(MARKER):
+            # AllReduce (psum)
+            out = jax.lax.psum(c, axis_name="device")
+            return out
+
+    mesh = create_mesh(SHARDING_STRATEGY)
+    # Input to AllReduce is the output of Matmul, which is C [M, N]
+    # In gemm_all_reduce, C is effectively replicated or sharded depending on strategy, 
+    # but here SHARDING_STRATEGY is NO_SHARDING, so it's replicated?
+    # Actually, let's double check gemm_all_reduce out_sharding.
+    # out_sharding = get_out_sharding(SHARDING_STRATEGY) -> P(None, None) for NO_SHARDING
+    
+    # So the input to THIS function should match the output of the GEMM part in gemm_all_reduce
+    # In gemm_all_reduce:
+    # f(x,y): ... return out
+    # out_sharding is P(None, None).
+    
+    # But wait, inside gemm_all_reduce's `f`, `c = acc.astype(dtype)`.
+    # `c` is local to the device in shard_map terms if check_rep=False and in_specs are P(None, None).
+    # Yes, `gemm_all_reduce` uses `in_specs=(lhs_sharding.spec, rhs_sharding.spec)`.
+    # For NO_SHARDING, lhs_sharding is P(None, None), rhs is P(None, None).
+    # So `c` is [M, N] per device.
+    
+    # So here, we want input `c` to be P(None, None) per device.
+    
+    input_sharding = get_out_sharding(SHARDING_STRATEGY) # Reusing this as it matched C's distribution
+    out_sharding = get_out_sharding(SHARDING_STRATEGY)
+
+    jit_sharded_f = jax.jit(
+        shard_map(
+            f,
+            mesh,
+            in_specs=(input_sharding,),
+            out_specs=out_sharding,
+            check_rep=False,
+        )
+    )
+
+    # Shape of C
+    c_shape = (m, n)
+    c_dtype = dtype
+
+    key = jax.random.key(SEED)
+
+    def data_generator():
+        """Creates new random data on host and puts it on device."""
+        nonlocal key
+        key, key_c = jax.random.split(key, 2)
+
+        # Create random data on host
+        c_host = jax.random.normal(key_c, c_shape).astype(c_dtype)
+
+        # Put on device (HBM)
+        # We need to wrap input_sharding (which is a PartitionSpec) in NamedSharding
+        # because device_put needs to know the mesh.
+        named_input_sharding = jax.sharding.NamedSharding(mesh, input_sharding)
+        c_device = jax.device_put(c_host, named_input_sharding)
+
+        return (c_device,)
+
+    time_ms_list = iteration_timeit(
+        jit_sharded_f,
+        data_generator,
+        matrix_dim=f"{dtype_str}_{m}x{n}x{k}",
+        tries=num_runs,
+        task=f"all_reduce_only_{dtype_str}",
+        trace_dir=trace_dir,
+    )
+    return {
+        "time_ms_list": time_ms_list,
+    }
+
+
+def all_reduce_only_calculate_metrics(
+    m: int,
+    k: int,
+    n: int,
+    dtype: jnp.dtype,
+    time_ms_list: list[float],
+) -> Dict[str, Any]:
+    
+    # Calculate Bandwidth for AllReduce
+    # AllReduce moves Matrix C: M x N
+    matrix_c_size_bytes = m * n * dtype.dtype.itemsize
+    
+    # Use unified_bytes_metrics for bandwidth-bound operations
+    # We estimate total_bytes_all_devices assuming full replication or reduction over all devices
+    num_devices = jax.device_count()
+    total_bytes_all_devices = matrix_c_size_bytes * num_devices
+
+    metadata, metrics = unified_bytes_metrics(
+        m, n, time_ms_list,
+        total_bytes=matrix_c_size_bytes,
+        total_bytes_all_devices=total_bytes_all_devices,
+        dtype=dtype.dtype.name
+    )
+    
+    return metadata, metrics

Ironwood/src/run_benchmark.py

@@ -62,6 +62,8 @@
     "gemm_throttling": "benchmark_gemm_throttling.gemm_throttling",
     "gemm": "benchmark_gemm.gemm",
     "gemm_all_reduce": "benchmark_gemm_all_reduce.gemm_all_reduce",
+    "gemm_only": "benchmark_gemm_all_reduce.gemm_only",
+    "all_reduce_only": "benchmark_gemm_all_reduce.all_reduce_only",
     "gemm_accum": "benchmark_gemm.gemm_accum",
     "quantization": "benchmark_compute.quantization",
     "transpose_quantization": "benchmark_compute.transpose_quantization",