upadte

Author: Binyang2014

python/mscclpp_benchmark/bench_collective.py

@@ -69,7 +69,6 @@ class DTypeSpec:
     cupy_dtype: Any
     mscclpp_dtype: Any
     accum_dtype: Any | None = None
-    accum_name: str | None = None
     fp8_format: str | None = None
 
 
@@ -126,7 +125,6 @@ def _parse_dtype(dtype_name: str) -> DTypeSpec:
             cp.uint8,
             mscclpp.DataType.float8_e4m3fn,
             accum_dtype=mscclpp.DataType.float16,
-            accum_name="float16",
             fp8_format="e4m3fn",
         )
     if normalized in {"float8_e4m3fnuz", "fp8_e4m3fnuz"}:
@@ -135,7 +133,6 @@ def _parse_dtype(dtype_name: str) -> DTypeSpec:
             cp.uint8,
             mscclpp.DataType.float8_e4m3fnuz,
             accum_dtype=mscclpp.DataType.float16,
-            accum_name="float16",
             fp8_format="e4m3fnuz",
         )
     if normalized in {"float8_e4m3b15", "fp8_e4m3b15"}:
@@ -144,7 +141,6 @@ def _parse_dtype(dtype_name: str) -> DTypeSpec:
             cp.uint8,
             mscclpp.DataType.float8_e4m3b15,
             accum_dtype=mscclpp.DataType.float32,
-            accum_name="float32",
             fp8_format="e4m3b15",
         )
     raise ValueError(
@@ -161,13 +157,10 @@ def _with_accum_type(dtype_spec: DTypeSpec, accum_type: str | None) -> DTypeSpec
     normalized = accum_type.strip().lower().replace("-", "_")
     if normalized in {"native", "same", "auto"}:
         accum_dtype = dtype_spec.mscclpp_dtype
-        accum_name = dtype_spec.name
     elif normalized in {"float16", "fp16", "half"}:
         accum_dtype = mscclpp.DataType.float16
-        accum_name = "float16"
     elif normalized in {"float32", "fp32", "float"}:
         accum_dtype = mscclpp.DataType.float32
-        accum_name = "float32"
     else:
         raise ValueError(f"Unsupported accum type {accum_type!r}; use native, float16, or float32")
 
@@ -176,11 +169,21 @@ def _with_accum_type(dtype_spec: DTypeSpec, accum_type: str | None) -> DTypeSpec
         cupy_dtype=dtype_spec.cupy_dtype,
         mscclpp_dtype=dtype_spec.mscclpp_dtype,
         accum_dtype=accum_dtype,
-        accum_name=accum_name,
         fp8_format=dtype_spec.fp8_format,
     )
 
 
+def _accum_dtype_name(dtype_spec: DTypeSpec) -> str:
+    accum_dtype = dtype_spec.accum_dtype
+    if accum_dtype is None or accum_dtype == dtype_spec.mscclpp_dtype:
+        return dtype_spec.name
+    if accum_dtype == _mscclpp().DataType.float16:
+        return "float16"
+    if accum_dtype == _mscclpp().DataType.float32:
+        return "float32"
+    raise ValueError(f"Unsupported accumulation data type: {accum_dtype}")
+
+
 def _human_size(size: int) -> str:
     value = float(size)
     for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
@@ -590,7 +593,7 @@ def main(argv: list[str] | None = None) -> None:
                 continue
             if args.autotune:
                 dtype = dtype_spec.name if args.collective == _ALLREDUCE else None
-                accum = (dtype_spec.accum_name or dtype_spec.name) if dtype is not None else None
+                accum = _accum_dtype_name(dtype_spec) if dtype is not None else None
                 config_store.upsert(
                     hardware_profile,
                     args.collective,