diff --git a/examples/mlperf/gpt_oss_20b/.dockerignore b/examples/mlperf/gpt_oss_20b/.dockerignore new file mode 100644 index 000000000..e0def13df --- /dev/null +++ b/examples/mlperf/gpt_oss_20b/.dockerignore @@ -0,0 +1,5 @@ +** +!Dockerfile.runtime-v26.3 +!Dockerfile.runtime-v26.5 +!prewarm_attention.py +!aiter_hd64_asm_override.py diff --git a/examples/mlperf/gpt_oss_20b/Dockerfile.runtime-v26.3 b/examples/mlperf/gpt_oss_20b/Dockerfile.runtime-v26.3 new file mode 100644 index 000000000..ed63cfd9b --- /dev/null +++ b/examples/mlperf/gpt_oss_20b/Dockerfile.runtime-v26.3 @@ -0,0 +1,88 @@ +# syntax=docker/dockerfile:1.7 + +ARG BASE_IMAGE=rocm/primus:v26.3@sha256:1a02b74a94d82131f3f119a94ccbed45bb5b4cd77a0616cb4e6a29e64a048482 +FROM ${BASE_IMAGE} + +ARG TRITON_REF=09500db9f0fe66fd176d1f080e2017b37e7e995d +ARG PRIMUS_TURBO_REF=fd8634f48b553029a58a7c0dbb03ba133fde3519 +ARG TE_REF=5235bae2cc683a0ad4bf15221c746ab3c1e229e7 +ARG AITER_REF=c4b33df03faae1c4e470420d950f8a9589e9634d +ARG FWD_ATTN_ASM_REF=53d3dadc3f3b0ac35ae536f2d1d7864a3e07ba22 +ARG BWD_ATTN_ASM_REF=9b9fb6444f3fee388617f62432c3faea74079377 +ARG BWD_ATTN_SYMBOL=_ZN5aiter43fmha_bwd_hd64_bf16_causal_a16_rtz_recompileE +ARG BWD_ATTN_SLOT=bwd_hd64_bf16_causal_a16_rtz.co +ARG MAX_JOBS=96 + +ENV MLPERF_RUNTIME_SERIES=v26.3 \ + MLPERF_ENABLE_FWD_ATTN_ASM=1 \ + FMHA_HD64_ASM_CO=/opt/venv/lib/python3.12/site-packages/aiter_hd64_asm_fwd_d64_opt128.co \ + FMHA_HD64_ASM_LOG=0 + +WORKDIR /workspace/deps + +# The v26.3 base ships Triton 3.6, while this Turbo branch requires the 3.7 API. +RUN python3 -m pip install --upgrade --force-reinstall --no-deps flydsl==0.2.4 && \ + git clone https://github.com/triton-lang/triton.git && \ + cd triton && git checkout "${TRITON_REF}" && \ + python3 -m pip install ninja cmake && \ + MAX_JOBS="${MAX_JOBS}" python3 -m pip install \ + --no-build-isolation --no-deps -v . && \ + cd .. && rm -rf triton + +RUN git clone --recursive https://github.com/AMD-AGI/Primus-Turbo.git && \ + cd Primus-Turbo && git checkout "${PRIMUS_TURBO_REF}" && \ + git submodule update --init --recursive && \ + python3 -m pip install -r requirements.txt && \ + PRIMUS_TURBO_FRAMEWORK=PYTORCH \ + GPU_ARCHS=gfx950 \ + MAX_JOBS="${MAX_JOBS}" \ + python3 -m pip install --no-build-isolation --no-deps -v . && \ + cd .. && rm -rf Primus-Turbo + +RUN git clone https://github.com/mawad-amd/fwd-attn-asm.git && \ + cd fwd-attn-asm && git checkout "${FWD_ATTN_ASM_REF}" && \ + amdclang -x assembler -target amdgcn-amd-amdhsa -mcpu=gfx950 \ + -o /tmp/fwd_attn.co kernels/fwd_d64_opt128.s && \ + cd .. && \ + git clone https://github.com/mawad-amd/bwd-attn-asm.git && \ + cd bwd-attn-asm && git checkout "${BWD_ATTN_ASM_REF}" && \ + amdclang -x assembler -target amdgcn-amd-amdhsa -mcpu=gfx950 \ + -o /tmp/bwd_attn.co kernels/bwd_d64_v3_causal_opt_16x32.s && \ + cd .. && rm -rf fwd-attn-asm bwd-attn-asm + +RUN git clone --recursive https://github.com/ROCm/TransformerEngine.git && \ + cd TransformerEngine && git checkout "${TE_REF}" && \ + git submodule update --init --recursive && \ + AITER_DIR=3rdparty/QoLA/3rdparty/aiter && \ + git -C "${AITER_DIR}" fetch origin "${AITER_REF}" && \ + git -C "${AITER_DIR}" checkout "${AITER_REF}" && \ + git -C "${AITER_DIR}" submodule update --init --recursive && \ + sed -i "s|^aiter_commit = .*|aiter_commit = \"${AITER_REF}\"|" \ + transformer_engine/common/ck_fused_attn/qola_manifest.toml && \ + install -m 0644 /tmp/bwd_attn.co \ + "${AITER_DIR}/hsa/gfx950/fmha_v3_bwd/${BWD_ATTN_SLOT}" && \ + python3 -m pip uninstall -y \ + transformer_engine transformer_engine_rocm7 transformer_engine_rocm_torch && \ + NVTE_FUSED_ATTN_AOTRITON=0 \ + NVTE_CK_FUSED_ATTN_PATH="" \ + NVTE_BUILD_MAX_JOBS="${MAX_JOBS}" \ + NVTE_FRAMEWORK=pytorch \ + NVTE_ROCM_ARCH=gfx950 \ + NVTE_USE_HIPBLASLT=1 \ + PYTORCH_ROCM_ARCH=gfx950 \ + CU_NUM=304 \ + NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD=1 \ + MAX_JOBS="${MAX_JOBS}" \ + python3 -m pip install --no-build-isolation --no-deps -v . && \ + cd .. && rm -rf TransformerEngine + +COPY aiter_hd64_asm_override.py \ + /opt/venv/lib/python3.12/site-packages/aiter_hd64_asm_override.py +COPY prewarm_attention.py /opt/mlperf-gpt-oss-20b/prewarm_attention.py +RUN printf 'import aiter_hd64_asm_override\n' \ + > /opt/venv/lib/python3.12/site-packages/aiter_hd64_asm_override.pth && \ + install -m 0644 /tmp/fwd_attn.co \ + /opt/venv/lib/python3.12/site-packages/aiter_hd64_asm_fwd_d64_opt128.co && \ + rm -f /tmp/fwd_attn.co /tmp/bwd_attn.co + +WORKDIR /workspace diff --git a/examples/mlperf/gpt_oss_20b/Dockerfile.runtime-v26.5 b/examples/mlperf/gpt_oss_20b/Dockerfile.runtime-v26.5 new file mode 100644 index 000000000..bcfc129fd --- /dev/null +++ b/examples/mlperf/gpt_oss_20b/Dockerfile.runtime-v26.5 @@ -0,0 +1,89 @@ +# syntax=docker/dockerfile:1.7 + +ARG BASE_IMAGE=rocm/primus:v26.5@sha256:3040bf42974d791dd42de2e36b3c919a00869a5754cfc57a06b96d004c55eed1 +FROM ${BASE_IMAGE} + +ARG PRIMUS_TURBO_REF=fd8634f48b553029a58a7c0dbb03ba133fde3519 +ARG TE_REF=a07e607f14a5330807ffdafeeb6224f2d7dffacc +ARG AITER_REF=d32b0cb62ecc32bb1a858e8437d58eb9b3856af6 +ARG FWD_ATTN_ASM_REF=53d3dadc3f3b0ac35ae536f2d1d7864a3e07ba22 +ARG BWD_ATTN_ASM_REF=9b9fb6444f3fee388617f62432c3faea74079377 +ARG BWD_ATTN_SOURCE_SYMBOL=_ZN5aiter43fmha_bwd_hd64_bf16_causal_a16_rtz_recompileE +ARG BWD_ATTN_SYMBOL=_ZN5aiter44fmha_bwd_hd64_bf16_causal_a16_rtne_recompileE +ARG BWD_ATTN_SLOT=bwd_hd64_bf16_causal_a16_rtne.co +ARG MAX_JOBS=96 + +ENV MLPERF_RUNTIME_SERIES=v26.5 \ + MLPERF_ENABLE_FWD_ATTN_ASM=1 \ + FMHA_HD64_ASM_CO=/opt/venv/lib/python3.12/site-packages/aiter_hd64_asm_fwd_d64_opt128.co \ + FMHA_HD64_ASM_LOG=0 + +WORKDIR /workspace/deps + +# The v26.5 base already contains the tested Triton 3.7 compiler. +RUN git clone --recursive https://github.com/AMD-AGI/Primus-Turbo.git && \ + cd Primus-Turbo && git checkout "${PRIMUS_TURBO_REF}" && \ + git submodule update --init --recursive && \ + python3 -m pip install -r requirements.txt && \ + PRIMUS_TURBO_FRAMEWORK=PYTORCH \ + GPU_ARCHS=gfx950 \ + MAX_JOBS="${MAX_JOBS}" \ + python3 -m pip install --no-build-isolation --no-deps -v . && \ + cd .. && rm -rf Primus-Turbo + +# Compile the hand-tuned forward and numerically validated backward payloads. +RUN git clone https://github.com/mawad-amd/fwd-attn-asm.git && \ + cd fwd-attn-asm && git checkout "${FWD_ATTN_ASM_REF}" && \ + amdclang -x assembler -target amdgcn-amd-amdhsa -mcpu=gfx950 \ + -o /tmp/fwd_attn.co kernels/fwd_d64_opt128.s && \ + cd .. && \ + git clone https://github.com/mawad-amd/bwd-attn-asm.git && \ + cd bwd-attn-asm && git checkout "${BWD_ATTN_ASM_REF}" && \ + sed "s/${BWD_ATTN_SOURCE_SYMBOL}/${BWD_ATTN_SYMBOL}/g" \ + kernels/bwd_d64_v3_causal_opt_16x32.s > /tmp/bwd_attn.s && \ + amdclang -x assembler -target amdgcn-amd-amdhsa -mcpu=gfx950 \ + -o /tmp/bwd_attn.co /tmp/bwd_attn.s && \ + cd .. && rm -rf fwd-attn-asm bwd-attn-asm /tmp/bwd_attn.s + +# The base image already carries the matching TE/AITER Python stack. Rebuild +# TE only because the validated backward ASM must be embedded in libmha_bwd. +RUN git clone --recursive https://github.com/ROCm/TransformerEngine.git && \ + cd TransformerEngine && git checkout "${TE_REF}" && \ + git submodule update --init --recursive && \ + sed -i "s|^aiter_commit = .*|aiter_commit = \"${AITER_REF}\"|" \ + transformer_engine/common/ck_fused_attn/qola_manifest.toml && \ + PYTHONPATH="$PWD/3rdparty/QoLA:${PYTHONPATH}" \ + python3 -m qola.cli checkout \ + --manifest transformer_engine/common/ck_fused_attn/qola_manifest.toml \ + --aiter-root /workspace/deps/te-aiter && \ + install -m 0644 /tmp/bwd_attn.co \ + "/workspace/deps/te-aiter/hsa/gfx950/fmha_v3_bwd/${BWD_ATTN_SLOT}" && \ + python3 -m pip uninstall -y \ + transformer_engine transformer_engine_rocm7 transformer_engine_rocm_torch && \ + NVTE_FUSED_ATTN_AOTRITON=0 \ + NVTE_FUSED_ATTN_CK=1 \ + NVTE_CK_JIT=1 \ + NVTE_CK_FUSED_ATTN_PATH="" \ + NVTE_AITER_SOURCE_DIR=/workspace/deps/te-aiter \ + NVTE_BUILD_MAX_JOBS="${MAX_JOBS}" \ + NVTE_FRAMEWORK=pytorch \ + NVTE_ROCM_ARCH=gfx950 \ + NVTE_USE_HIPBLASLT=1 \ + PYTORCH_ROCM_ARCH=gfx950 \ + CU_NUM=304 \ + NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD=1 \ + CMAKE_BUILD_PARALLEL_LEVEL="${MAX_JOBS}" \ + MAX_JOBS="${MAX_JOBS}" \ + python3 -m pip install --no-build-isolation --no-deps -v . && \ + cd .. && rm -rf TransformerEngine te-aiter + +COPY aiter_hd64_asm_override.py \ + /opt/venv/lib/python3.12/site-packages/aiter_hd64_asm_override.py +COPY prewarm_attention.py /opt/mlperf-gpt-oss-20b/prewarm_attention.py +RUN printf 'import aiter_hd64_asm_override\n' \ + > /opt/venv/lib/python3.12/site-packages/aiter_hd64_asm_override.pth && \ + install -m 0644 /tmp/fwd_attn.co \ + /opt/venv/lib/python3.12/site-packages/aiter_hd64_asm_fwd_d64_opt128.co && \ + rm -f /tmp/fwd_attn.co /tmp/bwd_attn.co + +WORKDIR /workspace diff --git a/examples/mlperf/gpt_oss_20b/README.md b/examples/mlperf/gpt_oss_20b/README.md index 8f29f31a0..0d979ce69 100644 --- a/examples/mlperf/gpt_oss_20b/README.md +++ b/examples/mlperf/gpt_oss_20b/README.md @@ -1,64 +1,77 @@ -# GPT-OSS-20B Pretraining Benchmark +# GPT-OSS-20B MLPerf pretraining -GPT-OSS 20B (Mixture of Experts) +GPT-OSS 20B on one MI355X node with 8 GPUs and global batch size 32. +## Build -## Setup - -### Start Docker Image +The Dockerfile builds the complete runtime, including the GPT-OSS +Primus-Turbo test branch, TransformerEngine, and the attention ASM kernels. +The v26.5 image reuses the base image's Triton 3.7 compiler; v26.3 upgrades its +older base Triton for compatibility with the same Turbo branch. ```bash -docker run -it --device /dev/dri --device /dev/kfd --device /dev/infiniband --network host --ipc host --group-add video --cap-add SYS_PTRACE --security-opt seccomp=unconfined --privileged -v $HOME:$HOME --shm-size 128G --name primus_training_env rocm/primus:v26.5 - -cd /workspace/Primus +cd examples/mlperf/gpt_oss_20b +docker build --network host \ + -f Dockerfile.runtime-v26.5 \ + -t primus:gpt-oss-20b-mlperf-v26.5 . ``` +Use `Dockerfile.runtime-v26.3` and a v26.3 tag for the compatibility stack. +Push a shared tag with `docker push `. -### Configuration - -This benchmark trains a 20B parameter GPT model with Mixture of Experts (MoE) architecture using the Primus framework on AMD GPUs. - -**Key Features:** -- 20B parameter MoE model -- Expert Parallelism (EP=8) -- FP8 hybrid precision training -- Primus Turbo optimizations (DeepEP, sync-free MoE) - -## Key Files - -- `configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml` - Model and training config - - Update `train_data_path` and `train_data_path` to your local downloaded location -- `config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh` - System config and env vars - - Update `PRIMUS_PATH` to clone Primus Repo - - Update `EXP`to `/examples/mlperf/configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml` -- `run_and_time.sh` - Run script - -### Data - -Download preprocessed C4 dataset: +## Data ```bash mkdir -p /data/gpt_oss_20b cd /data/gpt_oss_20b - -# Download training and validation data bash <(curl -s https://raw.githubusercontent.com/mlcommons/r2-downloader/refs/heads/main/mlc-r2-downloader.sh) \ - -d data https://training.mlcommons-storage.org/metadata/llama-3-1-8b-preprocessed-c4-dataset.uri + -d data \ + https://training.mlcommons-storage.org/metadata/llama-3-1-8b-preprocessed-c4-dataset.uri ``` -After download, you should see files with the following naming conventions: -- Training: `c4-train.en_6_text_document.bin` and `.idx` -- Validation: `c4-validation-91205-samples.en_text_document.bin` and `.idx` +Training uses the `c4-train.en_6_text_document` prefix and validation uses +`c4-validation-91205-samples.en_text_document`. + +## Run -The data directory is approximately **80 GB** and model directory is approximately **30 GB**. +```bash +docker run -it --rm \ + --privileged --network host --ipc host --shm-size 128g \ + --cap-add SYS_PTRACE --security-opt seccomp=unconfined \ + --device /dev/dri --device /dev/kfd --device /dev/infiniband \ + -v /path/to/Primus:/workspace/Primus \ + -v /path/to/data:/data \ + -v /path/to/model:/model \ + -v /path/to/results:/results \ + primus:gpt-oss-20b-mlperf-v26.5 bash +``` -### How to run +Inside the container: ```bash -export HF_TOKEN= +cd /workspace/Primus/examples/mlperf/gpt_oss_20b source config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh -bash run_and_time.sh +./run_and_time.sh ``` -## Notes -- `log_interval: 99999999` suppresses regular Primus logs +That is the complete long-run entry. The config is the single source of +submission defaults: MLPerf trainer, 1.2M iteration ceiling, 128-step warmup, +FP8 Triton grouped GEMM, fused wgrad accumulation, and disabled profiling. +Short diagnostics and backend ablations should override environment variables +outside the checked-in submission config. + +## v26.5 attention prewarm + +The v26.5 TE stack lazily compiles two attention variants. Starting eight +torchrun ranks against an empty cache can race while writing the same blobs. +`run_and_time.sh` therefore runs `prewarm_attention.py` once before timing; the +helper only populates the sliding-window and full-attention cache entries. +The v26.3 TE/AITER stack does not exhibit this cache race, so the prewarm is +skipped automatically for v26.3. + +## Key files + +- `Dockerfile.runtime-v26.3`, `Dockerfile.runtime-v26.5`: complete runtime builds +- `config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh`: submission defaults +- `run_and_time.sh`: benchmark entry +- `prewarm_attention.py`: v26.5-only attention cache prewarm diff --git a/examples/mlperf/gpt_oss_20b/aiter_hd64_asm_override.py b/examples/mlperf/gpt_oss_20b/aiter_hd64_asm_override.py new file mode 100644 index 000000000..ffec7f060 --- /dev/null +++ b/examples/mlperf/gpt_oss_20b/aiter_hd64_asm_override.py @@ -0,0 +1,812 @@ +"""Route eligible TransformerEngine FMHA forward calls to the pinned gfx950 +HD64 BF16 assembly kernel. + +The module is imported at Python startup through ``aiter_hd64_asm_override.pth`` +and remains inactive unless ``MLPERF_ENABLE_FWD_ATTN_ASM=1``. +""" + +from __future__ import annotations + +import ctypes +import inspect +import logging +import math +import os +import struct +import sys +from typing import Tuple + +logger = logging.getLogger("fwd_attn_asm_override") +# Per-dispatch logging is opt-in only: it fires once per attention call, so an +# inherited INFO root level (verbose runs) must not switch it on. +_DISPATCH_LOG = os.environ.get("FMHA_HD64_ASM_LOG", "0") == "1" +if _DISPATCH_LOG: + logging.basicConfig(level=logging.INFO) + logger.setLevel(logging.INFO) + + +_ENABLED = os.environ.get("MLPERF_ENABLE_FWD_ATTN_ASM", "0") == "1" +_AITER_ROPE_ENABLED = os.environ.get("NVTE_USE_AITER_ROPE", "0") == "1" +_DEFAULT_CO_PATH = os.path.join( + os.path.dirname(__file__), + "aiter_hd64_asm_fwd_d64_opt128.co", +) +_CO_PATH = os.environ.get("FMHA_HD64_ASM_CO", _DEFAULT_CO_PATH) +_KERNEL_NAME = b"fmha_fwd_d64_bf16_causal" + +# Tile shape baked into the kernel (BlockFmhaPipelineQRKSVSAsync<128,64,...>). +_BLOCK_M = 128 +_LDS_BYTES = 13056 +_BLOCK_THREADS = 256 + +_HIP_LIB = None +_CO_DATA: bytes | None = None +# HIP modules are bound to a device's primary context, so each rank/device +# needs its own handle. +_KFUNC_BY_DEV: dict = {} +_KMODULE_BY_DEV: dict = {} +_DISPATCH_COUNT = 0 + +# CK aux tensors are needed by TE's backward wrapper. Capture a compatible +# template on the first eligible call and then replace only its LSE tensor. +_AUX_CTX_TEMPLATES: dict = {} + + +def get_dispatch_count() -> int: + """Return successful hand-tuned kernel launches in this process.""" + return _DISPATCH_COUNT + + +def _ensure_hip_lib(): + global _HIP_LIB + if _HIP_LIB is not None: + return _HIP_LIB + # ROCm Python wheels ship runtime and development copies of libamdhip64. + # PyTorch is linked against the versioned runtime SONAME; opening the + # unversioned development symlink can create a second HIP runtime with a + # separate module/context registry, causing hipModuleGetFunction to return + # hipErrorNotFound even though the code object contains the symbol. + lib = ctypes.CDLL("libamdhip64.so.7") + lib.hipModuleLoadData.restype = ctypes.c_int + lib.hipModuleLoadData.argtypes = [ctypes.c_void_p, ctypes.c_void_p] + lib.hipModuleGetFunction.restype = ctypes.c_int + lib.hipModuleGetFunction.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_char_p, + ] + lib.hipModuleLaunchKernel.restype = ctypes.c_int + lib.hipModuleLaunchKernel.argtypes = [ + ctypes.c_void_p, + ctypes.c_uint, + ctypes.c_uint, + ctypes.c_uint, + ctypes.c_uint, + ctypes.c_uint, + ctypes.c_uint, + ctypes.c_uint, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ] + _HIP_LIB = lib + return _HIP_LIB + + +def _try_load_kernel() -> bool: + """Stage the code object; bind it to a HIP context on first launch.""" + global _CO_DATA + try: + _ensure_hip_lib() + if _CO_DATA is None: + with open(_CO_PATH, "rb") as file: + _CO_DATA = file.read() + return True + except Exception as error: # noqa: BLE001 + logger.warning( + "could not stage hand-tuned hd64 kernel path=%s exists=%s: %r", + _CO_PATH, + os.path.exists(_CO_PATH), + error, + ) + return False + + +def _get_kfunc_for_device(device) -> ctypes.c_void_p: + import torch + + dev_id = ( + device.index if hasattr(device, "index") and device.index is not None else torch.cuda.current_device() + ) + cached = _KFUNC_BY_DEV.get(dev_id) + if cached is not None: + return cached + + hip = _ensure_hip_lib() + if _CO_DATA is None: + with open(_CO_PATH, "rb") as file: + globals()["_CO_DATA"] = file.read() + + with torch.cuda.device(dev_id): + module = ctypes.c_void_p() + rc = hip.hipModuleLoadData( + ctypes.byref(module), + ctypes.create_string_buffer(_CO_DATA), + ) + if rc != 0: + raise RuntimeError(f"hipModuleLoadData failed for {_CO_PATH} on device {dev_id} " f"(rc={rc})") + func = ctypes.c_void_p() + rc = hip.hipModuleGetFunction( + ctypes.byref(func), + module, + _KERNEL_NAME, + ) + if rc != 0: + raise RuntimeError(f"hipModuleGetFunction failed on device {dev_id} (rc={rc})") + + _KMODULE_BY_DEV[dev_id] = module + _KFUNC_BY_DEV[dev_id] = func + logger.info( + "loaded hand-tuned hd64 kernel from %s on device %d", + _CO_PATH, + dev_id, + ) + return func + + +def _is_gfx950(device) -> bool: + import torch + + try: + return torch.cuda.get_device_properties(device).gcnArchName.startswith("gfx950") + except Exception: + return False + + +def _ck_window_args(window_size, attn_mask_type: str) -> Tuple[int, int]: + if window_size is None: + window_left, window_right = -1, -1 + else: + window_left, window_right = int(window_size[0]), int(window_size[1]) + if "causal" in (attn_mask_type or ""): + window_right = 0 + return window_left, window_right + + +def _eligible( + *, + max_seqlen_q, + max_seqlen_kv, + q, + k, + v, + attn_scale, + attn_bias_type, + attn_mask_type, + softmax_type, + window_size, + bottom_right_diagonal, + qkv_layout, + dropout, + attn_bias, + softmax_offset, + s_quantizer, + o_quantizer, + fp8, +) -> bool: + import torch + + if not _ENABLED or fp8: + return False + if q.dtype != torch.bfloat16 or k.dtype != torch.bfloat16 or v.dtype != torch.bfloat16: + return False + # TE 2.15 passes quantizer objects through its BF16 fused-attention API even + # when fp8=False. They are unused by this BF16 output path and must not make + # an otherwise eligible call fall back to CK. + if attn_bias is not None or softmax_offset is not None: + return False + if attn_bias_type != "no_bias" or softmax_type != "vanilla": + return False + if bottom_right_diagonal not in (None, False): + return False + if qkv_layout not in ("bshd_bshd_bshd", "sbhd_sbhd_sbhd"): + return False + if q.dim() != 4 or k.dim() != 4 or v.dim() != 4: + return False + if q.shape[-1] != 64 or v.shape[-1] != 64: + return False + if "causal" not in (attn_mask_type or ""): + return False + if dropout != 0.0: + return False + if not _is_gfx950(q.device): + return False + + if qkv_layout.startswith("bshd"): + batch, sequence_length, query_heads, _ = q.shape + key_value_sequence_length = k.shape[1] + key_value_heads = k.shape[2] + else: + sequence_length, batch, query_heads, _ = q.shape + key_value_sequence_length = k.shape[0] + key_value_heads = k.shape[2] + + if sequence_length != key_value_sequence_length: + return False + if max_seqlen_q != sequence_length or max_seqlen_kv != sequence_length: + return False + if q.shape[:2] != k.shape[:2] or k.shape[:2] != v.shape[:2]: + return False + if query_heads % key_value_heads != 0: + return False + if sequence_length % _BLOCK_M != 0: + return False + # TE uses 0.0 as a sentinel for the default 1/sqrt(head_dim) scale. + if attn_scale not in (None, 0.0) and not math.isclose( + float(attn_scale), + 1.0 / math.sqrt(q.shape[-1]), + rel_tol=1e-6, + abs_tol=0.0, + ): + return False + if batch <= 0: + return False + return True + + +def _launch( + q, + k, + v, + qkv_layout: str, + *, + attn_scale: float, + attn_mask_type: str, + window_size, + lse_out, +): + import torch + + if qkv_layout.startswith("bshd"): + batch, sequence_length, query_heads, head_dim = q.shape + key_value_heads = k.shape[2] + stride_q_s, stride_q_h, stride_q_b = ( + q.stride(1), + q.stride(2), + q.stride(0), + ) + stride_k_s, stride_k_h, stride_k_b = ( + k.stride(1), + k.stride(2), + k.stride(0), + ) + stride_v_s, stride_v_h, stride_v_b = ( + v.stride(1), + v.stride(2), + v.stride(0), + ) + else: + sequence_length, batch, query_heads, head_dim = q.shape + key_value_heads = k.shape[2] + stride_q_s, stride_q_h, stride_q_b = ( + q.stride(0), + q.stride(2), + q.stride(1), + ) + stride_k_s, stride_k_h, stride_k_b = ( + k.stride(0), + k.stride(2), + k.stride(1), + ) + stride_v_s, stride_v_h, stride_v_b = ( + v.stride(0), + v.stride(2), + v.stride(1), + ) + + output = torch.empty_like(q) + if qkv_layout.startswith("bshd"): + stride_o_s, stride_o_h, stride_o_b = ( + output.stride(1), + output.stride(2), + output.stride(0), + ) + else: + stride_o_s, stride_o_h, stride_o_b = ( + output.stride(0), + output.stride(2), + output.stride(1), + ) + + # The assembly kernel expects the CK log2(e)-scaled convention. + del attn_scale + scale_s = (1.0 / math.sqrt(head_dim)) * math.log2(math.e) + window_left, window_right = _ck_window_args( + window_size, + attn_mask_type, + ) + + kargs = struct.pack( + " 4 else None, + args[4].stride() if len(args) > 4 else None, + args[5].stride() if len(args) > 5 else None, + args[6].stride() if len(args) > 6 else None, + args[7].stride() if len(args) > 7 else None, + args[8].stride() if len(args) > 8 else None, + ) + try: + return original_fused_attn_bwd(*args, **kwargs) + except RuntimeError: + aux = args[11] if len(args) > 11 else kwargs.get("aux_ctx_tensors") + logger.error( + "fused-attn-bwd rejected config: max_q=%s max_kv=%s " + "q=%s backend=%s aux=%s layout=%s mask=%s window=%s " + "bottom_right=%s deterministic=%s", + args[0] if len(args) > 0 else None, + args[1] if len(args) > 1 else None, + getattr(args[4], "shape", None) if len(args) > 4 else None, + args[12] if len(args) > 12 else None, + [ + (tuple(t.shape), str(t.dtype)) + for t in (aux or []) + ], + args[21] if len(args) > 21 else None, + args[23] if len(args) > 23 else None, + args[25] if len(args) > 25 else None, + args[26] if len(args) > 26 else None, + args[27] if len(args) > 27 else None, + ) + raise + + diagnostic_fused_attn_bwd._fwd_attn_asm_bwd_diagnostic = True + fused_attn_module.fused_attn_bwd = diagnostic_fused_attn_bwd + + try: + from transformer_engine.pytorch.attention.dot_product_attention import backends + + backends.fused_attn_fwd = patched_fused_attn_fwd + backends.fused_attn_bwd = diagnostic_fused_attn_bwd + except Exception: + pass + + logger.info("fused_attn_fwd patched " "(D=64 BF16 [SWA-]causal -> hand-tuned hd64 kernel)") + return True + + +def _install_aiter_rope_override(rope_module): + """Restore the AITER RoPE route removed by newer ROCm TE revisions.""" + if not _AITER_ROPE_ENABLED: + return False + + fused_rope = rope_module.FusedRoPEFunc + if getattr(fused_rope, "_mlperf_aiter_rope_patched", False): + return True + # Older TE revisions already provide the same dispatch natively. + if hasattr(fused_rope, "_can_use_aiter"): + return True + + from aiter.ops.rope import rope_bwd as aiter_rope_bwd + from aiter.ops.rope import rope_fwd as aiter_rope_fwd + + original_forward = fused_rope.forward + original_backward = fused_rope.backward + + def aiter_aware_forward( + ctx, + tensor, + freqs, + start_positions=None, + tensor_format="sbhd", + interleaved=False, + cu_seqlens=None, + cp_size=1, + cp_rank=0, + ): + use_aiter = ( + tensor_format == "sbhd" + and not interleaved + and cu_seqlens is None + and cp_size == 1 + and start_positions is None + ) + ctx._mlperf_use_aiter_rope = use_aiter + if not use_aiter: + return original_forward( + ctx, + tensor, + freqs, + start_positions, + tensor_format, + interleaved, + cu_seqlens, + cp_size, + cp_rank, + ) + + if freqs.dtype != torch.float32: + freqs = freqs.float() + output = aiter_rope_fwd(tensor, freqs, 0, False, False) + ctx.save_for_backward(freqs, cu_seqlens, start_positions) + return output + + def aiter_aware_backward(ctx, grad_output): + if not getattr(ctx, "_mlperf_use_aiter_rope", False): + return original_backward(ctx, grad_output) + freqs, _, _ = ctx.saved_tensors + grad_input = aiter_rope_bwd(grad_output, freqs, 0, False, False) + return grad_input, None, None, None, None, None, None, None, None + + import torch + + fused_rope.forward = staticmethod(aiter_aware_forward) + fused_rope.backward = staticmethod(aiter_aware_backward) + fused_rope._mlperf_aiter_rope_patched = True + logger.info("restored AITER fused RoPE dispatch for this TE revision") + return True + + +class _DeferredInstaller: + _TARGETS = { + "transformer_engine.pytorch.attention.rope", + "transformer_engine.pytorch.cpp_extensions.fused_attn", + "transformer_engine.pytorch.cpp_extensions", + } + + def find_spec(self, fullname, path, target=None): + try: + if fullname not in self._TARGETS: + return None + for finder in sys.meta_path: + if finder is self: + continue + try: + spec = finder.find_spec(fullname, path, target) + except (AttributeError, ImportError): + spec = None + if spec is None: + continue + original_loader = spec.loader + + class _WrappedLoader: + def create_module(self, module_spec): + if hasattr(original_loader, "create_module"): + return original_loader.create_module(module_spec) + return None + + def exec_module(self, module): + original_loader.exec_module(module) + if fullname == "transformer_engine.pytorch.cpp_extensions.fused_attn": + try: + _install_fused_attn_override() + except Exception as error: # noqa: BLE001 + logger.warning( + "deferred install failed: %r", + error, + ) + elif fullname == "transformer_engine.pytorch.attention.rope": + try: + _install_aiter_rope_override(module) + except Exception as error: # noqa: BLE001 + logger.warning( + "deferred AITER RoPE install failed: %r", + error, + ) + + spec.loader = _WrappedLoader() + return spec + return None + except Exception as error: # noqa: BLE001 + logger.warning( + "fwd-attn-asm find_spec error for %s: %r", + fullname, + error, + ) + return None + + +def _register_deferred_install(): + if any(isinstance(finder, _DeferredInstaller) for finder in sys.meta_path): + return + sys.meta_path.insert(0, _DeferredInstaller()) + logger.info("deferred installer registered; will patch on TE load") + + +if _ENABLED or _AITER_ROPE_ENABLED: + try: + _register_deferred_install() + except Exception as error: # noqa: BLE001 + logger.warning( + "fwd-attn-asm deferred install failed at startup: %r", + error, + ) diff --git a/examples/mlperf/gpt_oss_20b/config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh b/examples/mlperf/gpt_oss_20b/config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh index 4309dbd4e..7b5248081 100644 --- a/examples/mlperf/gpt_oss_20b/config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh +++ b/examples/mlperf/gpt_oss_20b/config_MI355X_1x8x1_tp1pp1ep1_gbs32.sh @@ -48,7 +48,10 @@ export PRIMUS_EP=1 # ----------------------------------------------------------------------------- # Primus Configuration # ----------------------------------------------------------------------------- -export PRIMUS_TURBO_GROUPED_GEMM_BACKEND=TRITON +export PRIMUS_TURBO_GROUPED_GEMM_BACKEND="${PRIMUS_TURBO_GROUPED_GEMM_BACKEND:-triton}" +export PRIMUS_TURBO_GEMM_BACKEND=triton +export PRIMUS_TURBO_FUSED_WGRAD_ACCUM="${PRIMUS_TURBO_FUSED_WGRAD_ACCUM:-1}" +export PRIMUS_NUM_WORKERS="${PRIMUS_NUM_WORKERS:-2}" export PRIMUS_GRAD_REDUCE_IN_BF16=true export USE_TURBO_RMS_NORM=true @@ -84,7 +87,14 @@ export NCCL_CHECKS_DISABLE=1 # ----------------------------------------------------------------------------- export USE_HIPBLASLT=1 export TORCH_BLAS_PREFER_HIPBLASLT=1 -export HIPBLASLT_TUNING_OVERRIDE_FILE=${PRIMUS_PATH}/examples/mlperf/gpt_oss_20b/tune_gemm_results.txt +HIPBLASLT_TUNING_OVERRIDE_FILE="${PRIMUS_PATH}/examples/mlperf/gpt_oss_20b/tune_gemm_results-${MLPERF_RUNTIME_SERIES:-v26.3}.txt" +if [ -f "${HIPBLASLT_TUNING_OVERRIDE_FILE}" ]; then + export HIPBLASLT_TUNING_OVERRIDE_FILE +else + # hipBLASLt solution indices are runtime-specific. Never feed the v26.3 + # catalog to v26.5 while its own tuning table has not been generated. + unset HIPBLASLT_TUNING_OVERRIDE_FILE +fi # ----------------------------------------------------------------------------- # NVTE — FP8 & Cast Transpose @@ -100,19 +110,27 @@ export NVTE_FLASH_ATTN=0 # Disable FlashAttention so FusedAtten export NVTE_CK_USES_FWD_V3=1 # Globally on; aiter selects v3 vs CK-tile internally export NVTE_CK_USES_BWD_V3=1 # Globally on; aiter selects v3 vs CK-tile internally export NVTE_USE_AITER_ROPE=1 # Route RoPE through aiter's fused kernel instead of TE's own CK kernel -export NVTE_FMHA_USE_BSHD=0 # Native SBHD path (aiter c4b33df0 supports it; skips Megatron's SBHD↔BSHD shim transposes) -export NVTE_CK_IS_V3_ATOMIC_FP32=1 # use atomic fp32 kernels for now. atomic fp16 kernels resulting in numerics issues. -export NVTE_CK_HOW_V3_BF16_CVT=2 # 0=RTNE, 1=RTNA, 2=RTZ - -# fwd-attn-asm: route eligible (D=64 BF16 [SWA-]causal) fused_attn_fwd calls -# to the hand-tuned gfx950 kernel staged into site-packages by the Dockerfile. -# Set to 0 to disable. FMHA_HD64_ASM_LOG=1 prints one line per dispatch. -export MLPERF_ENABLE_FWD_ATTN_ASM=1 -export FMHA_HD64_ASM_LOG=0 - -# bwd-attn-asm is build-time only — TE's QoLA build embeds aiter's bwd `.co` -# into te_libmha_bwd.so at pip-install. Toggle with Docker `--build-arg -# BWD_ATTN_ASM_ENABLE=0` (default 1) at image build time. +export NVTE_FMHA_USE_BSHD=0 # Use the native SBHD path provided by the v26.5 AITER stack +export NVTE_CK_HOW_V3_BF16_CVT=2 # 0=RTNE, 1=RTNA, 2=RTZ (gfx942 selector; gfx950 is fixed) + +# Route eligible D=64 BF16 causal forward calls to the hand-tuned gfx950 +# kernel. FMHA_HD64_ASM_LOG=1 prints one line per successful launch. +export MLPERF_ENABLE_FWD_ATTN_ASM="${MLPERF_ENABLE_FWD_ATTN_ASM:-1}" +export FMHA_HD64_ASM_LOG="${FMHA_HD64_ASM_LOG:-0}" + +# The RTZ custom backward code object is registered under the RTNE-named slot +# that v26.5 AITER actually selects on gfx950. Unlike the bundled v26.5 a16 +# kernel that overflowed at sequence length 8192, this injected kernel passed a +# 200-step FP8 C4 convergence run. Set the flag to 0 for the a32 PSSK fallback. +export MLPERF_ENABLE_BWD_ATTN_ASM="${MLPERF_ENABLE_BWD_ATTN_ASM:-1}" +if [ "${MLPERF_ENABLE_BWD_ATTN_ASM}" = "1" ]; then + export NVTE_CK_IS_V3_ATOMIC_FP32=0 +elif [ "${MLPERF_ENABLE_BWD_ATTN_ASM}" = "0" ]; then + export NVTE_CK_IS_V3_ATOMIC_FP32=1 +else + echo "MLPERF_ENABLE_BWD_ATTN_ASM must be 0 or 1" >&2 + return 2 2>/dev/null || exit 2 +fi # ----------------------------------------------------------------------------- # NVTE — Debug @@ -156,16 +174,29 @@ export SYNTH_WARMUP_STEPS=3 # ----------------------------------------------------------------------------- # Skip sort_chunks_by_idxs when the per-local-expert index is an identity # permutation (fires at EP=1/TP=1). Set to 0 to run the original path; useful -# for A/B measurements. See patches/megatron_moe_skip_identity_sort.patch. +# for A/B measurements. Implemented by skip_identity_sort_patches.py. export MOE_SKIP_IDENTITY_SORT=1 # ----------------------------------------------------------------------------- # DDP Parameter All-Gather (SDMA) # ----------------------------------------------------------------------------- -export ENABLE_SDMA_ALLGATHER=1 +# v26.5 turns five SDMA workspace barriers into ~5 ms waits each. A same-node +# A/B with the latest Turbo main measured 1004 ms/step with SDMA and 988 +# ms/step with RCCL, matching the v26.3 control at 989 ms/step. Keep SDMA +# opt-in on v26.5 until its barrier regression is fixed; retain the validated +# v26.3 default. +if [ "${MLPERF_RUNTIME_SERIES:-v26.5}" = "v26.3" ]; then + DEFAULT_ENABLE_SDMA_ALLGATHER=1 +else + DEFAULT_ENABLE_SDMA_ALLGATHER=0 +fi +export ENABLE_SDMA_ALLGATHER="${ENABLE_SDMA_ALLGATHER:-${DEFAULT_ENABLE_SDMA_ALLGATHER}}" # Optional: cap the per-call peer-copy stream count. Default is # min(world_size-1, 8); lower values reduce SDMA / memory-system pressure. # export MEGATRON_SDMA_PEER_COPY_STREAMS=8 +# When SDMA is explicitly enabled, keep the source-patch behavior: the first +# two parameter bucket groups use RCCL and later groups use SDMA. +export MEGATRON_SDMA_RCCL_FALLBACK_BUCKETS=${MEGATRON_SDMA_RCCL_FALLBACK_BUCKETS:-2} # ----------------------------------------------------------------------------- # Run-log verbosity diff --git a/examples/mlperf/gpt_oss_20b/configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml b/examples/mlperf/gpt_oss_20b/configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml index 9074e4419..00efb9ebb 100644 --- a/examples/mlperf/gpt_oss_20b/configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml +++ b/examples/mlperf/gpt_oss_20b/configs/MI355/gpt_oss_20B-FP8-mlperf-pretrain.yaml @@ -1,7 +1,7 @@ work_group: ${TEAM:amd} user_name: ${USER:root} exp_name: ${EXP_NAME:gpt_oss_20b} -workspace: ./output +workspace: ${PRIMUS_WORKSPACE:./output} modules: pre_trainer: @@ -13,7 +13,7 @@ modules: overrides: # Activate the migrated MLPerf pretrain trainer (mllog + MLPerf hooks). - stage: mlperf_pretrain + stage: ${PRIMUS_STAGE:mlperf_pretrain} # tokenizer tokenizer_type: Llama3Tokenizer @@ -39,7 +39,7 @@ modules: # log wandb_project: "Primus_GPT_OSS_20B" stderr_sink_level: ERROR - log_interval: 999999 + log_interval: ${LOG_INTERVAL:999999} # debug # moe_router_force_load_balancing: true @@ -51,10 +51,10 @@ modules: use_pytorch_profiler: ${PRIMUS_PROFILE:False} profile_step_end: ${PRIMUS_PROFILE_STEP_END:32} profile_step_start: ${PRIMUS_PROFILE_STEP_START:16} - profile_ranks: [0,1,2,3,4,5,6,7] + profile_ranks: [0] # enable fp8 training - fp8: e4m3 + fp8: ${PRIMUS_FP8:e4m3} fp8_recipe: tensorwise clip_grad: 1.0 # Gradient clipping (already default, but explicit) check_for_nan_in_loss_and_grad: false @@ -150,6 +150,8 @@ modules: moe_router_load_balancing_type: none moe_router_num_groups: null moe_router_padding_for_fp8: false + # Keep routing exact. PrimusTurbo consumes the ragged expert token counts + # directly, so no fake routes or TE multi_padding/unpadding are needed. moe_router_padding_for_quantization: false moe_router_pre_softmax: false moe_router_score_function: softmax @@ -183,14 +185,13 @@ modules: # Turbo enable_primus_turbo: true use_turbo_attention: false - use_turbo_grouped_gemm: false + use_turbo_grouped_gemm: true + use_turbo_ragged_grouped_gemm: ${USE_TURBO_RAGGED_GROUPED_GEMM:true} use_turbo_rms_norm: ${USE_TURBO_RMS_NORM:true} + # Keeps the fused layernorm+QKV site's norm on Turbo and its GEMM on + # hipBLASLt, which is where each backend is faster on these shapes. + use_turbo_norm_te_linear: ${USE_TURBO_NORM_TE_LINEAR:true} use_turbo_fused_act_with_probs : true - # Pad tokens-per-expert so the fp8 grouped GEMM path skips the buggy - # quantization_padding branch in PrimusGroupedMLP.forward (experts.py:97-109), - # which yields NaN with recompute. Not auto-enabled here because - # turbo_sync_free_moe_stage=0 (it is only auto-set for sync-free stages 1-3). - use_turbo_permute_padding: true # deepep use_turbo_deepep: false @@ -209,8 +210,8 @@ modules: cross_entropy_loss_fusion: true # tensorboard logging, set 'disable_tensorboard: false' to enable tensorboard logging - disable_tensorboard: true - tensorboard_dir: /workspace/code/tensorboard + disable_tensorboard: ${PRIMUS_DISABLE_TENSORBOARD:true} + tensorboard_dir: ${PRIMUS_WORKSPACE:./output}/tensorboard tensorboard_log_interval: 1 tensorboard_queue_size: 1000 log_timers_to_tensorboard: true diff --git a/examples/mlperf/gpt_oss_20b/prewarm_attention.py b/examples/mlperf/gpt_oss_20b/prewarm_attention.py new file mode 100644 index 000000000..a505256b2 --- /dev/null +++ b/examples/mlperf/gpt_oss_20b/prewarm_attention.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +"""Populate the two GPT-OSS attention JIT cache entries before torchrun.""" + +import argparse + +import torch +from transformer_engine.pytorch.attention import DotProductAttention + +import aiter_hd64_asm_override + + +def prewarm(sequence_length: int, micro_batch_size: int, window_left: int) -> None: + attention = DotProductAttention( + num_attention_heads=64, + kv_channels=64, + num_gqa_groups=8, + attention_dropout=0.0, + qkv_format="sbhd", + attn_mask_type="causal", + window_size=(window_left, 0), + ).cuda() + shapes = ( + (sequence_length, micro_batch_size, 64, 64), + (sequence_length, micro_batch_size, 8, 64), + (sequence_length, micro_batch_size, 8, 64), + ) + query, key, value = ( + torch.randn(shape, device="cuda", dtype=torch.bfloat16, requires_grad=True) + for shape in shapes + ) + output = attention(query, key, value) + if isinstance(output, tuple): + output = output[0] + output.float().square().mean().backward() + torch.cuda.synchronize() + print(f"attention_prewarm=PASS window_left={window_left}") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--sequence-length", type=int, default=8192) + parser.add_argument("--micro-batch-size", type=int, default=4) + parser.add_argument("--window-left", type=int, action="append") + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise RuntimeError("ROCm GPU is not available") + torch.cuda.set_device(0) + torch.manual_seed(1234) + + windows = args.window_left or [128, -1] + for window_left in windows: + prewarm(args.sequence_length, args.micro_batch_size, window_left) + if aiter_hd64_asm_override.get_dispatch_count() < len(windows): + raise RuntimeError("Forward ASM override was not dispatched") + + +if __name__ == "__main__": + main() diff --git a/examples/mlperf/gpt_oss_20b/run_and_time.sh b/examples/mlperf/gpt_oss_20b/run_and_time.sh index fc07e6fe1..1ef73b1ad 100755 --- a/examples/mlperf/gpt_oss_20b/run_and_time.sh +++ b/examples/mlperf/gpt_oss_20b/run_and_time.sh @@ -6,16 +6,23 @@ set -e mkdir -p /results cd "${PRIMUS_PATH}/examples/mlperf/gpt_oss_20b" +TRAIN_LOG_FILE="${TRAIN_LOG_FILE:-train.mlperfpretrain.exp.log}" -# Under multi-node SLURM (run_with_docker_slurm.sh), inherit rendezvous + node -# sizing from SLURM env so we can scale to N nodes without editing the config -# file. Single-node SLURM jobs (NNODES=1) fall through to the config defaults -# so torchrun doesn't try to do c10d rdzv against MASTER_ADDR=localhost. +# Under a multi-node scheduler wrapper, inherit rendezvous + node sizing from +# SLURM so the same benchmark config scales without edits. Single-node jobs +# fall through to the config defaults. if [[ -n "${SLURM_NNODES:-}" && "${SLURM_NNODES}" -gt 1 ]]; then NNODES="${SLURM_NNODES}" NODE_RANK="${SLURM_NODEID:-0}" fi +# TE 2.15 lazily compiles CK attention blobs. Populate both GPT-OSS windows +# once before torchrun so eight ranks do not race while writing the cache. +if [ "${MLPERF_RUNTIME_SERIES:-v26.3}" = "v26.5" ] \ + && [ "${MLPERF_SKIP_ATTENTION_PREWARM:-0}" != "1" ]; then + python3 /opt/mlperf-gpt-oss-20b/prewarm_attention.py +fi + echo "============================================" echo "MLPerf GPT-OSS-20B Training" echo "============================================" @@ -38,7 +45,7 @@ set +e "${PRIMUS_PATH}/primus-cli" direct -- \ train pretrain \ --config "${EXP}" \ - 2>&1 | tee train.mlperfpretrain.exp.log + 2>&1 | tee "${TRAIN_LOG_FILE}" ret_code=${PIPESTATUS[0]} set -e diff --git a/examples/mlperf/gpt_oss_20b/tune_gemm_results.txt b/examples/mlperf/gpt_oss_20b/tune_gemm_results-v26.3.txt similarity index 100% rename from examples/mlperf/gpt_oss_20b/tune_gemm_results.txt rename to examples/mlperf/gpt_oss_20b/tune_gemm_results-v26.3.txt diff --git a/examples/mlperf/gpt_oss_20b/tune_gemm_results-v26.5.txt b/examples/mlperf/gpt_oss_20b/tune_gemm_results-v26.5.txt new file mode 100644 index 000000000..4e3711ba3 --- /dev/null +++ b/examples/mlperf/gpt_oss_20b/tune_gemm_results-v26.5.txt @@ -0,0 +1,41 @@ +Git Version: fa9cdd18 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + T,N,0,1,32,32768,2880,1,2880,92160,0,2880,94371840,32,1048576,32,1048576,bf16_r,bf16_r,f32_r,f32_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,f32_r,512,1,0,182699,5440.59,33.0588,24641,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + N,N,0,1,2880,32768,32,1,2880,92160,0,32,1048576,2880,94371840,2880,94371840,f32_r,f32_r,f32_r,f32_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,f32_r,512,1,0,68985.3,4064.01,87.552,73624,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + N,T,0,1,4096,2880,32768,1,4096,134217728,0,2880,94371840,4096,11796480,4096,11796480,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.4418e+06,835.047,536.202,12468,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + N,T,0,1,4096,2880,32768,1,4096,134217728,1,2880,94371840,4096,11796480,4096,11796480,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.4024e+06,891.947,551.265,11292,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + N,T,0,1,2880,32,32768,1,2880,94371840,0,32,1048576,2880,92160,2880,92160,f32_r,f32_r,f32_r,f32_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,f32_r,512,1,0,86893.9,5119.03,69.5077,72656,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + N,T,0,1,2880,5120,32768,1,2880,94371840,0,5120,167772160,2880,14745600,2880,14745600,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.42455e+06,760.278,678.366,11258,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + N,T,0,1,2880,5120,32768,1,2880,94371840,1,5120,167772160,2880,14745600,2880,14745600,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.38626e+06,818.644,697.102,12554,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + T,N,0,1,2880,32768,4096,1,4096,11796480,0,4096,134217728,2880,94371840,2880,94371840,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.43674e+06,832.117,538.09,19081,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + N,N,0,1,4096,32768,2880,1,4096,11796480,0,2880,94371840,4096,134217728,4096,134217728,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.47155e+06,852.281,525.36,14593,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + N,N,0,1,2880,32768,5120,1,2880,14745600,0,5120,167772160,2880,94371840,2880,94371840,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.46461e+06,781.659,659.811,14734,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + T,N,0,1,5120,32768,2880,1,2880,14745600,0,2880,94371840,5120,167772160,5120,167772160,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.53294e+06,818.125,630.401,21406,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + N,N,0,1,2880,32768,128256,1,2880,369377280,0,128256,4202692608,2880,94371840,2880,94371840,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.53955e+06,552.788,15723.8,13283,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + N,T,0,1,2880,128256,32768,1,2880,94371840,1,128256,4202692608,2880,369377280,2880,369377280,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.49105e+06,620.132,16235.2,12111,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + T,N,0,1,128256,32768,2880,1,2880,369377280,0,2880,94371840,128256,4202692608,128256,4202692608,bf16_r,bf16_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,1.46089e+06,524.544,16570.4,22416,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + T,N,0,1,4096,2880,32768,1,32768,134217728,0,32768,94371840,4096,11796480,4096,11796480,f8_r,f8_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.79595e+06,849.399,276.505,37702,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + T,N,0,1,2880,32768,4096,1,4096,11796480,0,4096,134217728,2880,94371840,2880,94371840,f8_r,f8_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.81692e+06,1135.99,274.447,38550,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + T,N,0,1,2880,5120,32768,1,32768,94371840,0,32768,167772160,2880,14745600,2880,14745600,f8_r,f8_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.97207e+06,835.326,325.15,36173,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + T,N,0,1,2880,32768,5120,1,5120,14745600,0,5120,167772160,2880,94371840,2880,94371840,f8_r,f8_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.87346e+06,1028.12,336.308,38461,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + T,N,0,1,4096,32768,2880,1,2880,11796480,0,2880,94371840,4096,134217728,4096,134217728,f8_r,f8_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.56755e+06,1158.67,301.102,36093,gfx950:sramecc+:xnack-,256 + transA,transB,grouped_gemm,batch_count,m,n,k,alpha,lda,stride_a,beta,ldb,stride_b,ldc,stride_c,ldd,stride_d,a_type,b_type,c_type,d_type,compute_type,scaleA,scaleB,scaleC,scaleD,amaxD,swizzle_a,swizzle_b,activation_type,bias_vector,bias_type,aux_type,rotating_buffer,flush,use_gpu_timer,hipblaslt-Gflops,hipblaslt-GB/s,us,solution_index,gcnArchName,CUs + T,N,0,1,5120,32768,2880,1,2880,14745600,0,2880,94371840,5120,167772160,5120,167772160,f8_r,f8_r,bf16_r,bf16_r,f32_r,0,0,0,0,0,0,0,none,0,f32_r,bf16_r,512,1,0,2.54105e+06,1088.93,380.303,36094,gfx950:sramecc+:xnack-,256 diff --git a/primus/backends/megatron/core/distributed/sdma_param_gather.py b/primus/backends/megatron/core/distributed/sdma_param_gather.py index 4f2f5ebb8..758e516d7 100644 --- a/primus/backends/megatron/core/distributed/sdma_param_gather.py +++ b/primus/backends/megatron/core/distributed/sdma_param_gather.py @@ -22,7 +22,7 @@ Activation: ``ENABLE_SDMA_ALLGATHER=1`` (gated by the patch). When the required - Primus-Turbo / ``hip`` primitives are unavailable, every call falls back to + Primus-Turbo primitives are unavailable, every call falls back to ``torch.distributed.all_gather_into_tensor`` so behaviour is preserved. """ @@ -166,11 +166,11 @@ def all_gather_into_tensor_sdma( assert input_tensor.is_contiguous(), "SDMA all_gather_into_tensor requires contiguous input_tensor" try: - hip = importlib.import_module("hip").hip - get_amd_symm_mem_workspace = importlib.import_module( - "primus_turbo.pytorch.kernels.async_tp.amd_symmetric_memory" - ).get_amd_symm_mem_workspace - hip_check = importlib.import_module("primus_turbo.pytorch.kernels.async_tp.common_ops").hip_check + symm_mem_module = importlib.import_module("primus_turbo.pytorch.core.symm_mem") + hip_runtime_module = importlib.import_module("primus_turbo.pytorch.core.pyhip_runtime_wrapper") + get_symm_mem_workspace = symm_mem_module.get_symm_mem_workspace + hip_runtime = hip_runtime_module.get_hip_runtime_lib() + memcpy_comm_kind = hip_runtime_module.hipMemcpyKindEnum.hipMemcpyDeviceToDeviceNoCU except Exception: return _all_gather_into_tensor_waitable_fallback( output_tensor, input_tensor, group=group, async_op=async_op @@ -185,12 +185,11 @@ def all_gather_into_tensor_sdma( input_nbytes = input_tensor.nbytes output_flat = output_tensor.view(-1) - memcpy_comm_kind = getattr( - hip.hipMemcpyKind, "hipMemcpyDeviceToDeviceNoCU", hip.hipMemcpyKind.hipMemcpyDeviceToDevice - ) - # Allocate/reuse symmetric workspace and expose each rank's local shard buffer. - symm_mem = get_amd_symm_mem_workspace(group_name, min_size=max(input_nbytes, _SDMA_SYMM_MEM_MIN_BYTES)) + symm_mem = get_symm_mem_workspace( + group, + min_size=max(input_nbytes, _SDMA_SYMM_MEM_MIN_BYTES), + ) gather_buffers = [ symm_mem.get_buffer(r, input_tensor.shape, input_tensor.dtype) for r in range(world_size) ] @@ -211,14 +210,12 @@ def all_gather_into_tensor_sdma( # Peer copies may start after the publish barrier completes. barrier_event.record(runtime.comm_stream) # Copy local shard into its slot in the output. - hip_check( - hip.hipMemcpyAsync( - output_flat.data_ptr() + rank * input_nbytes, - input_tensor.data_ptr(), - input_nbytes, - memcpy_comm_kind, - runtime.comm_stream.cuda_stream, - ) + hip_runtime.hipMemcpyAsync( + output_flat.data_ptr() + rank * input_nbytes, + input_tensor.data_ptr(), + input_nbytes, + memcpy_comm_kind, + runtime.comm_stream.cuda_stream, ) for peer_idx, src_rank in enumerate(r for r in range(world_size) if r != rank): @@ -226,14 +223,12 @@ def all_gather_into_tensor_sdma( copy_stream = runtime.peer_copy_streams[peer_idx % len(runtime.peer_copy_streams)] with torch.cuda.stream(copy_stream): copy_stream.wait_event(barrier_event) - hip_check( - hip.hipMemcpyAsync( - output_flat.data_ptr() + src_rank * input_nbytes, - src_buf.data_ptr(), - input_nbytes, - memcpy_comm_kind, - copy_stream.cuda_stream, - ) + hip_runtime.hipMemcpyAsync( + output_flat.data_ptr() + src_rank * input_nbytes, + src_buf.data_ptr(), + input_nbytes, + memcpy_comm_kind, + copy_stream.cuda_stream, ) def _wait_impl(): diff --git a/primus/backends/megatron/core/extensions/primus_turbo.py b/primus/backends/megatron/core/extensions/primus_turbo.py index 05456b295..ffa81e9b4 100644 --- a/primus/backends/megatron/core/extensions/primus_turbo.py +++ b/primus/backends/megatron/core/extensions/primus_turbo.py @@ -1664,6 +1664,129 @@ def forward_internal(self, x, is_first_microbatch: bool = False): return out, None +def _make_primus_turbo_norm_te_column_parallel_linear(): + """Build a layernorm+linear module that keeps the GEMM on Transformer Engine. + + ``PrimusTurboLayerNormColumnParallelLinear`` above routes both the norm and + the GEMM through Turbo. On the GPT-OSS-20B shapes only the norm is worth + moving: Turbo's Triton rmsnorm is markedly cheaper than TE's ``general`` + kernels (measured 21.3 ms/step against 26.9 with the dgamma fix, and 38.5 + without it), while Turbo's dense FP8 GEMM is slower than hipBLASLt here + (69.6 ms/step against 59.7). + + TE fuses norm and GEMM inside one autograd function, so there is no way to + substitute just the norm. This composes ``PrimusTurboRMSNorm`` with + ``TEColumnParallelLinear`` instead, trading the fused launch for the faster + norm. Built in a factory so the TE import stays lazy. + """ + import torch.nn as nn + from megatron.core.extensions.transformer_engine import ( + TEColumnParallelLinear, + _get_extra_te_kwargs, + ) + + class PrimusTurboNormTEColumnParallelLinear(nn.Module): + """See ``_make_primus_turbo_norm_te_column_parallel_linear``.""" + + def __init__( + self, + input_size: int, + output_size: int, + *, + config: TransformerConfig, + init_method: Callable, + gather_output: bool, + bias: bool, + skip_bias_add: bool, + is_expert: bool, + skip_weight_param_allocation: bool = False, + tp_comm_buffer_name: Optional[str] = None, + tp_group: Optional[torch.distributed.ProcessGroup] = None, + stride: int = 1, + ): + super().__init__() + self.config = config + # Mirror the kwargs TELayerNormColumnParallelLinear feeds into + # te.pytorch.LayerNormLinear, so the norm behaves identically. + extra = _get_extra_te_kwargs(config) + self.layernorm = PrimusTurboRMSNorm( + normalized_shape=input_size, + eps=config.layernorm_epsilon, + sequence_parallel=config.sequence_parallel, + zero_centered_gamma=config.layernorm_zero_centered_gamma, + params_dtype=extra.get("params_dtype", torch.float32), + device=extra.get("device", torch.cuda.current_device()), + ) + self.linear = TEColumnParallelLinear( + input_size=input_size, + output_size=output_size, + config=config, + init_method=init_method, + gather_output=gather_output, + bias=bias, + skip_bias_add=skip_bias_add, + is_expert=is_expert, + skip_weight_param_allocation=skip_weight_param_allocation, + tp_comm_buffer_name=tp_comm_buffer_name, + tp_group=tp_group, + ) + + # Callers reach for these names on the fused TE module directly + # (te_op_fuser, modelopt state-dict hooks, checkpoint conversion), so + # expose them rather than forcing every caller to know the layout. + @property + def layer_norm_weight(self): + return self.layernorm.weight + + @property + def normalization(self): + return self.config.normalization + + @property + def weight(self): + return self.linear.weight + + @property + def bias(self): + return getattr(self.linear, "bias", None) + + @property + def in_features(self): + return self.linear.in_features + + @property + def out_features(self): + return self.linear.out_features + + def forward(self, x): + return self.linear(self.layernorm(x)) + + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): + """Flatten to the names a checkpoint from the fused class carries.""" + sd = {} + sd.update( + self.linear.sharded_state_dict( + prefix=prefix, sharded_offsets=sharded_offsets, metadata=metadata + ) + ) + norm_sd = self.layernorm.state_dict(prefix="", keep_vars=True) + sd.update( + make_sharded_tensors_for_checkpoint( + {"layer_norm_weight": norm_sd["weight"]}, prefix, {}, sharded_offsets + ) + ) + return sd + + def __repr__(self): + return ( + f"{type(self).__name__}(in_features={self.in_features}, " + f"out_features={self.out_features}, " + f"norm={type(self.layernorm).__name__})" + ) + + return PrimusTurboNormTEColumnParallelLinear + + def fused_bias_act_with_probs( intermediate_parallel: torch.Tensor, bias_parallel: torch.Tensor, diff --git a/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py b/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py index 48aedcd43..4bb006e5a 100644 --- a/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py +++ b/primus/backends/megatron/core/extensions/transformer_engine_spec_provider.py @@ -43,6 +43,7 @@ PrimusTurboLinear, PrimusTurboRowParallelGroupedLinear, PrimusTurboRowParallelLinear, + _make_primus_turbo_norm_te_column_parallel_linear, ) except (ImportError, ModuleNotFoundError): PrimusTurboAttention = None @@ -52,6 +53,7 @@ PrimusTurboLinear = None PrimusTurboRowParallelGroupedLinear = None PrimusTurboRowParallelLinear = None + _make_primus_turbo_norm_te_column_parallel_linear = None _LEGACY_GROUPED_MLP_CLS = None @@ -108,6 +110,7 @@ def _build_default_primus_args() -> SimpleNamespace: return SimpleNamespace( enable_primus_turbo=False, use_turbo_gemm=False, + use_turbo_norm_te_linear=False, use_turbo_attention=False, use_turbo_grouped_gemm=False, moe_use_legacy_grouped_gemm=False, @@ -190,14 +193,23 @@ def fuse_layernorm_and_linear(self) -> bool: return True def column_parallel_layer_norm_linear(self) -> Optional[type]: - """Which module for sequential layernorm and linear""" - return ( - _require_primus_turbo( + """Which module for sequential layernorm and linear. + + Three choices, because the norm and the GEMM want different backends on + these shapes: Turbo's rmsnorm beats TE's, while TE's hipBLASLt GEMM beats + Turbo's. ``use_turbo_norm_te_linear`` takes only the norm. + """ + if self.cfg.use_turbo_gemm: + return _require_primus_turbo( PrimusTurboLayerNormColumnParallelLinear, "layernorm column parallel linear" ) - if self.cfg.use_turbo_gemm - else TELayerNormColumnParallelLinear - ) + if getattr(self.cfg, "use_turbo_norm_te_linear", False): + factory = _require_primus_turbo( + _make_primus_turbo_norm_te_column_parallel_linear, + "turbo-norm layernorm column parallel linear", + ) + return factory() + return TELayerNormColumnParallelLinear def layer_norm(self, rms_norm: bool = False, for_qk: bool = False) -> type: """Which module to use for layer norm""" diff --git a/primus/backends/megatron/core/transformer/experts.py b/primus/backends/megatron/core/transformer/experts.py index 694db2137..469639a9a 100644 --- a/primus/backends/megatron/core/transformer/experts.py +++ b/primus/backends/megatron/core/transformer/experts.py @@ -46,6 +46,36 @@ def __init__( # NOTE: use_turbo_fused_act_with_probs is prioritized over use_te_activation_func and bias_activation_fusion self.use_turbo_fused_act_with_probs = args.use_turbo_fused_act_with_probs self.moe_router_padding_for_quantization = args.moe_router_padding_for_quantization + self.use_turbo_ragged_grouped_gemm = getattr(args, "use_turbo_ragged_grouped_gemm", False) + # PrimusTurbo's tensorwise FP8 grouped GEMM consumes the original GPU + # tokens_per_expert tensor, including non-aligned (ragged) group sizes. + # Keep TE's explicit zero-padding fallback for other recipes/backends. + + def _use_explicit_quantization_padding(self) -> bool: + """Whether this forward must pad expert groups before quantization.""" + if not (self.config.fp8 or self.config.fp4): + return False + if self.moe_router_padding_for_quantization: + return False + if not self.use_turbo_ragged_grouped_gemm: + return True + + # Check the active autocast state at forward time. The CLI recipe alone + # does not prove that this grouped linear is actually using Turbo FP8. + from primus.backends.megatron.core.extensions.primus_turbo import ( + PrimusTurboLowPrecisionGlobalStateManager, + ) + + if not PrimusTurboLowPrecisionGlobalStateManager.is_turbo_fp8_enabled(): + raise RuntimeError( + "use_turbo_ragged_grouped_gemm=True requires an active PrimusTurbo FP8 autocast context." + ) + quant_config = PrimusTurboLowPrecisionGlobalStateManager.get_turbo_quant_config() + if not quant_config.current_scaling(): + raise RuntimeError( + "use_turbo_ragged_grouped_gemm=True currently requires tensorwise dynamic FP8 scaling." + ) + return False def bias_act_func(self, intermediate_parallel, bias_parallel, permuted_probs): """ @@ -170,7 +200,8 @@ def forward( Return: output (torch.Tensor): The output of the local experts. """ - if not self.moe_router_padding_for_quantization and (self.config.fp8 or self.config.fp4): + use_explicit_quantization_padding = self._use_explicit_quantization_padding() + if use_explicit_quantization_padding: # NOTE: When moe_router_padding_for_quantization is true the token is padded. So we can skip the padding here to reduce cpu sync. tokens_per_expert_cpu: list[int] = tokens_per_expert.tolist() actual_tokens_per_expert_cpu: list[int] = tokens_per_expert_cpu @@ -218,7 +249,14 @@ def forward( ) else: with off_interface(self.offload_moe_act, fc1_output, "moe_act") as fc1_output: - bias_act_output = self.bias_act_func(fc1_output, bias_parallel, permuted_probs) + # Honor use_turbo_fused_act_with_probs independently of activation + # recomputation. The helper falls back to bias_act_func when disabled. + bias_act_output = self.bias_act_func_with_mask( + fc1_output, + bias_parallel, + permuted_probs, + tokens_per_expert, + ) output, output_bias = apply_module(self.linear_fc2)(bias_act_output, tokens_per_expert) if self.activation_recompute: self.activation_checkpoint.discard_output_and_register_recompute(output) @@ -230,7 +268,7 @@ def forward( output = self._apply_bias(output, output_bias, tokens_per_expert, permuted_probs) # upad and concat the output - if not self.moe_router_padding_for_quantization and (self.config.fp8 or self.config.fp4): + if use_explicit_quantization_padding: output = self.quantization_unpadding(output, actual_tokens_per_expert_cpu) output_bias = None diff --git a/primus/backends/megatron/patches/args/rocm_arg_validation.py b/primus/backends/megatron/patches/args/rocm_arg_validation.py index 95aff1b15..8fffcdbf8 100644 --- a/primus/backends/megatron/patches/args/rocm_arg_validation.py +++ b/primus/backends/megatron/patches/args/rocm_arg_validation.py @@ -97,6 +97,25 @@ def validate_fsdp2_optimizer_exclusivity(args) -> None: ) +def validate_turbo_ragged_grouped_gemm(args) -> None: + """Validate the no-padding PrimusTurbo grouped-GEMM path.""" + option = "use_turbo_ragged_grouped_gemm" + if not getattr(args, option, False): + return + if not getattr(args, "enable_primus_turbo", False) or not getattr(args, "use_turbo_grouped_gemm", False): + raise ValueError(f"{option}=True requires enable_primus_turbo=True and use_turbo_grouped_gemm=True.") + if ( + not getattr(args, "fp8", None) + or getattr(args, "fp8_recipe", None) != "tensorwise" + or getattr(args, "fp4", False) + ): + raise ValueError(f"{option}=True currently supports only tensorwise FP8.") + if getattr(args, "moe_router_padding_for_quantization", False): + raise ValueError( + "use_turbo_ragged_grouped_gemm=True requires moe_router_padding_for_quantization=False." + ) + + def validate_args_on_rocm(args): # Primus-Turbo auto-tuning use_turbo_autotune = getattr(args, "use_turbo_autotune", False) @@ -187,6 +206,8 @@ def validate_args_on_rocm(args): f"========== Enable Sync-Free MoE Stage {args.turbo_sync_free_moe_stage} (Auto-Enabled Options) ==========" ) + validate_turbo_ragged_grouped_gemm(args) + # turbo deepep if args.use_turbo_deepep: assert ( diff --git a/primus/backends/megatron/patches/parallelism/sdma_param_all_gather_patches.py b/primus/backends/megatron/patches/parallelism/sdma_param_all_gather_patches.py index f44875b78..f0aabc446 100644 --- a/primus/backends/megatron/patches/parallelism/sdma_param_all_gather_patches.py +++ b/primus/backends/megatron/patches/parallelism/sdma_param_all_gather_patches.py @@ -16,8 +16,9 @@ 1. ``_ParamAndGradBucketGroup.start_param_sync`` -- the distributed-optimizer path is re-implemented to dispatch one all-gather per bucket through :func:`all_gather_into_tensor_sdma` (copy-engine) instead of the RCCL - ``_coalescing_manager`` group. The first two bucket groups (by gather - order) stay on the regular RCCL fallback. The layer-wise optimizer path is + ``_coalescing_manager`` group. By default, the first two bucket groups (by + gather order) stay on the regular RCCL fallback; this is tunable with + ``MEGATRON_SDMA_RCCL_FALLBACK_BUCKETS``. The layer-wise optimizer path is delegated unchanged to the original method. 2. ``DistributedDataParallel.__init__`` -- after construction, each bucket group is annotated with ``param_gather_order`` (reverse dispatch order, @@ -33,6 +34,7 @@ """ import os +import warnings import torch @@ -44,6 +46,19 @@ def _sdma_allgather_enabled(_ctx: PatchContext) -> bool: return os.environ.get("ENABLE_SDMA_ALLGATHER", "0") == "1" +def _get_rccl_fallback_bucket_count() -> int: + value = os.getenv("MEGATRON_SDMA_RCCL_FALLBACK_BUCKETS", "2") + try: + count = int(value) + except ValueError: + warnings.warn(f"Invalid MEGATRON_SDMA_RCCL_FALLBACK_BUCKETS={value!r}; using 2.") + return 2 + if count < 0: + warnings.warn(f"MEGATRON_SDMA_RCCL_FALLBACK_BUCKETS must be non-negative; got {count}, using 2.") + return 2 + return count + + def _make_start_param_sync(orig_start_param_sync): """Build a replacement ``start_param_sync`` for ``_ParamAndGradBucketGroup``.""" from megatron.core.distributed.param_and_grad_buffer import shard_buffer @@ -70,15 +85,17 @@ def start_param_sync(self, force_sync: bool = False): async_op = self.ddp_config.overlap_param_gather and not force_sync - # Keep the first two bucket groups (by gather order) on the regular - # RCCL all-gather; route the rest through SDMA. param_gather_order is - # assigned in the DDP __init__ wrapper below. + # Keep a configurable number of leading bucket groups on regular RCCL; + # route the rest through SDMA. param_gather_order is assigned in the + # DDP __init__ wrapper below. param_gather_order = getattr(self, "param_gather_order", None) enable_sdma = os.getenv("ENABLE_SDMA_ALLGATHER") == "1" + rccl_fallback_bucket_count = _get_rccl_fallback_bucket_count() + use_rccl_fallback = not enable_sdma or ( + param_gather_order is not None and param_gather_order < rccl_fallback_bucket_count + ) all_gather_func = ( - _all_gather_into_tensor_waitable_fallback - if (param_gather_order is not None and param_gather_order < 2) or not enable_sdma - else all_gather_into_tensor_sdma + _all_gather_into_tensor_waitable_fallback if use_rccl_fallback else all_gather_into_tensor_sdma ) param_gather_handles = [] @@ -117,7 +134,8 @@ def _make_wrapped_ddp_init(orig_init): def __init__(self, *args, **kwargs): orig_init(self, *args, **kwargs) # Mirror the source patch: number bucket groups in reverse dispatch - # order so start_param_sync can keep the first two on RCCL. + # order so start_param_sync can keep the configured leading groups on + # RCCL. for groups_attr in ("bucket_groups", "expert_parallel_bucket_groups"): groups = getattr(self, groups_attr, None) or [] for order, bucket_group in enumerate(reversed(groups)): diff --git a/primus/configs/modules/megatron/primus_turbo.yaml b/primus/configs/modules/megatron/primus_turbo.yaml index b1aefe04f..f7dbbdbd4 100644 --- a/primus/configs/modules/megatron/primus_turbo.yaml +++ b/primus/configs/modules/megatron/primus_turbo.yaml @@ -33,6 +33,9 @@ use_turbo_gemm: false # ===== Grouped MLP ===== # use turbo grouped gemm use_turbo_grouped_gemm: false +# let tensorwise FP8 Turbo grouped GEMM consume non-aligned expert token counts +# directly, without TE Fp8Padding/Fp8Unpadding +use_turbo_ragged_grouped_gemm: false # fused activation_with_probs to reduce redundant computation use_turbo_fused_act_with_probs: false @@ -55,3 +58,7 @@ use_turbo_mega_moe: false # ===== Layer Norm ===== # operator switch use_turbo_rms_norm: false +# Route the fused layernorm+QKV site through Turbo's rmsnorm while leaving its +# GEMM on Transformer Engine. use_turbo_gemm moves that GEMM to Turbo as well, +# which costs more than the norm saves on the GPT-OSS-20B shapes. +use_turbo_norm_te_linear: false diff --git a/tests/unit_tests/backends/megatron/test_rocm_arg_validation.py b/tests/unit_tests/backends/megatron/test_rocm_arg_validation.py new file mode 100644 index 000000000..2e9bb97a5 --- /dev/null +++ b/tests/unit_tests/backends/megatron/test_rocm_arg_validation.py @@ -0,0 +1,54 @@ +############################################################################### +# Copyright (c) 2026, Advanced Micro Devices, Inc. +# +# See LICENSE for license information. +############################################################################### + +from types import SimpleNamespace + +import pytest + +from primus.backends.megatron.patches.args.rocm_arg_validation import ( + validate_turbo_ragged_grouped_gemm, +) + + +def _ragged_args(**overrides): + values = { + "use_turbo_ragged_grouped_gemm": True, + "enable_primus_turbo": True, + "use_turbo_grouped_gemm": True, + "fp8": "e4m3", + "fp8_recipe": "tensorwise", + "fp4": False, + "moe_router_padding_for_quantization": False, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def test_turbo_ragged_grouped_gemm_accepts_tensorwise_fp8(): + validate_turbo_ragged_grouped_gemm(_ragged_args()) + + +def test_turbo_ragged_grouped_gemm_disabled_is_noop(): + validate_turbo_ragged_grouped_gemm(SimpleNamespace()) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"enable_primus_turbo": False}, "requires enable_primus_turbo"), + ({"use_turbo_grouped_gemm": False}, "requires enable_primus_turbo"), + ({"fp8": None}, "supports only tensorwise FP8"), + ({"fp8_recipe": "blockwise"}, "supports only tensorwise FP8"), + ({"fp4": True}, "supports only tensorwise FP8"), + ( + {"moe_router_padding_for_quantization": True}, + "requires moe_router_padding_for_quantization=False", + ), + ], +) +def test_turbo_ragged_grouped_gemm_rejects_unsupported_config(overrides, message): + with pytest.raises(ValueError, match=message): + validate_turbo_ragged_grouped_gemm(_ragged_args(**overrides))