diff --git a/.github/actions/build-pytorch-wheel/Dockerfile b/.github/actions/build-pytorch-wheel/Dockerfile index a858307ab4..05ac20db0d 100644 --- a/.github/actions/build-pytorch-wheel/Dockerfile +++ b/.github/actions/build-pytorch-wheel/Dockerfile @@ -41,9 +41,5 @@ RUN CUDA_MAJOR_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $1'}) && \ # Install PyTorch RUN export MATRIX_CUDA_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $1 $2'}) && \ export MATRIX_TORCH_VERSION=$(echo $TORCH_VERSION | awk -F \. {'print $1 "." $2'}) && \ - export TORCH_CUDA_VERSION=$(python -c "from os import environ as env; \ - minv = {'2.5': 118, '2.6': 118, '2.7': 118, '2.8': 126, '2.9': 126}[env['MATRIX_TORCH_VERSION']]; \ - maxv = {'2.5': 124, '2.6': 126, '2.7': 128, '2.8': 129, '2.9': 130}[env['MATRIX_TORCH_VERSION']]; \ - print(minv if int(env['MATRIX_CUDA_VERSION']) < 120 else maxv)" \ - ) && \ - pip install --no-cache-dir torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION} \ No newline at end of file + export TORCH_CUDA_VERSION=$(python -c "from os import environ as env; versions = {'2.5': (118, 124), '2.6': (118, 126), '2.7': (118, 128), '2.8': (126, 129), '2.9': (126, 130)}; minv, maxv = versions[env['MATRIX_TORCH_VERSION']]; print(minv if int(env['MATRIX_CUDA_VERSION']) < 120 else maxv)") && \ + pip install --no-cache-dir torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION} diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 2ef6d1893d..0f05dbc40a 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -1,66 +1,160 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. # A workflow to trigger TE build on GitHub - name: 'Build' on: pull_request: workflow_dispatch: +concurrency: + # Group by workflow name + PR number (for PRs) or ref (for branch/tag pushes) + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true jobs: + core: + name: 'Core' + runs-on: ubuntu-latest + container: + image: nvcr.io/nvidia/cuda:12.1.0-devel-ubuntu22.04 + options: --user root + steps: + - name: 'Dependencies' + run: | + apt-get update + apt-get install -y git python3.9 pip cudnn9-cuda-12 + pip install cmake==3.21.0 pybind11[global] ninja + - name: 'Checkout' + uses: actions/checkout@v3 + with: + submodules: recursive + - name: ccache + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad + - name: 'Build' + run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation . -v + env: + NVTE_FRAMEWORK: none + MAX_JOBS: 1 + SCCACHE_GHA_ENABLED: "true" + - name: 'Sanity check' + run: python3 -c "import transformer_engine" + working-directory: / pytorch: name: 'PyTorch' - runs-on: [ nv-8g-cicd-te ] - defaults: - run: - shell: bash - container: - image: harbor.baai.ac.cn/flagscale/cuda12.8.1-torch2.7.1-python3.10-te2.9:20260209 - ports: - - 80:80 - options: >- - --gpus all - --shm-size=500g - --privileged - --ipc=host - --ulimit memlock=-1 - --ulimit stack=67108864 - --ulimit nofile=65535:65535 - --user root - --pull never + runs-on: ubuntu-latest steps: - - name: Configure Git Safe Directory on Cuda - run: /usr/bin/git config --global safe.directory '*' + - name: Move /var/lib/docker/ + shell: bash -euxo pipefail {0} + run: sudo mv /var/lib/docker/ "${GITHUB_WORKSPACE}/docker" + + - name: Maximize build space + uses: easimon/maximize-build-space@c28619d8999a147d5e09c1199f84ff6af6ad5794 + with: + root-reserve-mb: 5120 + temp-reserve-mb: 32 + swap-size-mb: 10240 + remove-dotnet: 'true' + remove-android: 'true' + remove-haskell: 'true' + remove-codeql: 'true' + build-mount-path: '/var/lib/docker/' + + - name: Restore /var/lib/docker/ + shell: bash -euxo pipefail {0} + run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/* /var/lib/docker" - name: 'Checkout' - uses: actions/checkout@v4 + uses: actions/checkout@v3 with: - fetch-depth: 0 submodules: recursive - set-safe-directory: true - - name: 'Setup Environment' + - name: Start named container run: | - source /opt/miniconda3/etc/profile.d/conda.sh - conda activate flagscale-train - echo "PATH=$PATH" >> $GITHUB_ENV + docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d nvcr.io/nvidia/cuda:12.8.0-devel-ubuntu22.04 sleep infinity - - name: 'Build' + - name: 'Dependencies' run: | - pip uninstall transformer_engine transformer_engine_torch -y || true - echo "GITHUB_WORKSPACE=$GITHUB_WORKSPACE" - cd $GITHUB_WORKSPACE - pip install nvdlfw-inspect - pip install expecttest - pip install . -v --no-deps --no-build-isolation + docker exec builder bash -c '\ + apt-get update && \ + apt-get install -y git python3.9 pip cudnn9-cuda-12 && \ + pip install cmake torch ninja pydantic importlib-metadata>=1.0 packaging pybind11 numpy einops onnxscript && \ + apt-get clean \ + ' + + - name: 'Build' + run: docker exec builder bash -c 'pip install --no-build-isolation . -v --no-deps' env: NVTE_FRAMEWORK: pytorch - TE_WITH_NCCL: '1' - NVTE_WITH_CUDA: '1' - CUDA_HOME: /usr/local/cuda-12.8 - NVCC: /usr/local/cuda-12.8/bin/nvcc + MAX_JOBS: 1 + - name: 'Sanity check' + run: docker exec builder bash -c 'python3 tests/pytorch/test_sanity_import.py' + jax: + name: 'JAX' + runs-on: ubuntu-latest + container: + image: ghcr.io/nvidia/jax:jax + options: --user root + steps: + - name: 'Dependencies' + run: pip install cmake==3.21.0 pybind11[global] + - name: 'Checkout' + uses: actions/checkout@v3 + with: + submodules: recursive + - name: ccache + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad + - name: 'Build' + run: | + NVTE_CCACHE_BIN=sccache NVTE_USE_CCACHE=1 pip install --no-build-isolation . -v + env: + NVTE_FRAMEWORK: jax + MAX_JOBS: 1 + SCCACHE_GHA_ENABLED: "true" + - name: 'Sanity check' + run: python3 tests/jax/test_sanity_import.py + all: + name: 'All' + runs-on: ubuntu-latest + steps: + - name: Move /var/lib/docker/ + shell: bash -euxo pipefail {0} + run: sudo mv /var/lib/docker/ "${GITHUB_WORKSPACE}/docker" + + - name: Maximize build space + uses: easimon/maximize-build-space@c28619d8999a147d5e09c1199f84ff6af6ad5794 + with: + root-reserve-mb: 5120 + temp-reserve-mb: 32 + swap-size-mb: 10240 + remove-dotnet: 'true' + remove-android: 'true' + remove-haskell: 'true' + remove-codeql: 'true' + build-mount-path: '/var/lib/docker/' + - name: Restore /var/lib/docker/ + shell: bash -euxo pipefail {0} + run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/* /var/lib/docker" + + - name: 'Checkout' + uses: actions/checkout@v3 + with: + submodules: recursive + + - name: Start named container + run: | + docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d ghcr.io/nvidia/jax:jax sleep infinity + + - name: 'Dependencies' + run: | + docker exec builder bash -c '\ + pip install cmake==3.21.0 pybind11[global] einops onnxscript && \ + pip install torch --no-cache-dir --index-url https://download.pytorch.org/whl/cu130 + ' + - name: 'Build' + run: docker exec builder bash -c 'pip install --no-cache-dir --no-build-isolation . -v --no-deps' + env: + NVTE_FRAMEWORK: all + MAX_JOBS: 1 - name: 'Sanity check' - run: - python3 tests/pytorch/test_sanity_import.py + run: docker exec builder bash -c 'python3 tests/pytorch/test_sanity_import.py && python3 tests/jax/test_sanity_import.py' diff --git a/.github/workflows/community_label.yml b/.github/workflows/community_label.yml new file mode 100644 index 0000000000..c0d31d2a45 --- /dev/null +++ b/.github/workflows/community_label.yml @@ -0,0 +1,63 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# A workflow to automatically label the contributions as community/org +name: Label community contributions + +on: + pull_request_target: + types: [opened, reopened, ready_for_review, synchronize] + +permissions: + contents: read + issues: write + pull-requests: write + +jobs: + label: + runs-on: ubuntu-latest + steps: + - uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 + with: + script: | + const pr = context.payload.pull_request; + const user = pr.user.login; + const association = pr.author_association; + + const communityLabel = "community-contribution"; + const orgLabel = "org-contribution"; + + let targetLabel = null; + + const isOrgMember = + association === "MEMBER" || association === "OWNER"; + + let permission = "none"; + + try { + const res = await github.rest.repos.getCollaboratorPermissionLevel({ + owner: context.repo.owner, + repo: context.repo.repo, + username: user, + }); + permission = res.data.permission; + } catch (e) { + if (e.status !== 404) throw e; + } + + const isCore = permission === "write" || permission === "admin"; + if (!isOrgMember) { + targetLabel = communityLabel; + } else { + targetLabel = orgLabel; + } + + if (!isCore) { + await github.rest.issues.addLabels({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: pr.number, + labels: [targetLabel], + }); + } diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 1d2fb272f8..016d2079d2 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -11,6 +11,8 @@ concurrency: # Group by workflow name + PR number (for PRs) or ref (for branch/tag pushes) group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} cancel-in-progress: true +permissions: + contents: read jobs: pytorch_cpplint: name: 'PyTorch C++' diff --git a/.gitignore b/.gitignore index f1cedd9955..878ceff9e4 100644 --- a/.gitignore +++ b/.gitignore @@ -48,6 +48,7 @@ artifacts/ transformer_engine/plugin/core/_build_config.py # Mac OS .DS_Store +.claude/ # Integration test outputs qa/L1_pytorch_mcore_integration/output/ *.distcp diff --git a/.gitmodules b/.gitmodules index 4b188d6bb1..495d8e3fe7 100644 --- a/.gitmodules +++ b/.gitmodules @@ -7,3 +7,6 @@ [submodule "3rdparty/cutlass"] path = 3rdparty/cutlass url = https://github.com/NVIDIA/cutlass.git +[submodule "3rdparty/nccl"] + path = 3rdparty/nccl + url = https://github.com/NVIDIA/nccl.git diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 76f476eb3f..601149916b 100755 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -40,7 +40,7 @@ repos: files: ^transformer_engine.*\.(c|cc|cxx|cpp|cu|cuh|h|hpp)$ - repo: https://github.com/netromdk/vermin - rev: c75aca72f4e85c6e47252139e8695f1c8b5f9ae3 + rev: b70ff9611a01a2bf2f702aa537d14e71e330edba hooks: - id: vermin args: ['-t=3.10-', '--violations'] diff --git a/3rdparty/nccl b/3rdparty/nccl new file mode 160000 index 0000000000..a6b5de08b6 --- /dev/null +++ b/3rdparty/nccl @@ -0,0 +1 @@ +Subproject commit a6b5de08b6af4f938cef541ae6e4d405632f89a4 diff --git a/CODEOWNERS b/CODEOWNERS new file mode 100644 index 0000000000..3087832fa4 --- /dev/null +++ b/CODEOWNERS @@ -0,0 +1,24 @@ +# IMPORTANT: +# This file is ONLY used to subscribe for notifications for PRs +# related to a specific file path. Approvals from people in this +# file are not required for merges. + +# C API +/transformer_engine/common/include/ @ptrendx + +# TE/JAX +/transformer_engine/jax/ @jberchtold-nvidia + +# TE/PyTorch +/transformer_engine/pytorch/ @ksivaman + +# te.ops API +/transformer_engine/pytorch/ops/ @timmoon10 + +# Quantization kernels +/transformer_engine/common/cast/ @Oleg-Goncharov + +# Attention +/transformer_engine/pytorch/attention/ @cyanguwa +/transformer_engine/common/fused_attn/ @cyanguwa +/transformer_engine/jax/cpp_extensions/attention.py @KshitijLakhani diff --git a/README.rst b/README.rst index 13f60bd72e..ae19e37e7f 100644 --- a/README.rst +++ b/README.rst @@ -5,17 +5,15 @@ |License| - -**TransformerEngine-FL is a fork of TransformerEngine that introduces a plugin-based architecture for supporting diverse AI chips, built on top of** `FlagOS `_, **a unified open-source AI system software stack.** - Transformer Engine ================== -`Quickstart <#examples>`_ | `Installation <#installation>`_ | `User Guide `_ | `Examples `_ | `FP8 Convergence <#fp8-convergence>`_ | `Integrations <#integrations>`_ | `Release notes `_ +`Quickstart <#examples>`_ | `Installation <#installation>`_ | `User Guide `_ | `Examples `_ | `Convergence <#convergence>`_ | `Integrations <#integrations>`_ | `Release notes `_ Latest News =========== +* [12/2025] `NVIDIA Nemotron 3: Efficient and Open Intelligence `_ - trained with NVFP4 on Transformer Engine * [11/2025] `NVIDIA Blackwell Architecture Sweeps MLPerf Training v5.1 Benchmarks `_ * [11/2025] `Scale Biology Transformer Models with PyTorch and NVIDIA BioNeMo Recipes `_ * [11/2025] `FP8 Training of Large-Scale RL Models `_ @@ -33,27 +31,26 @@ What is Transformer Engine? Transformer Engine (TE) is a library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit floating point (FP8) precision on Hopper, Ada, and Blackwell GPUs, to provide better -performance with lower memory utilization in both training and inference. TE provides a collection +performance with lower memory utilization in both training and inference. On Blackwell GPUs, TE also +supports MXFP8 (Microscaling FP8) and NVFP4 formats for even greater efficiency. TE provides a collection of highly optimized building blocks for popular Transformer architectures and an automatic mixed precision-like API that can be used seamlessly with your framework-specific code. TE also includes a framework agnostic C++ API that can be integrated with other deep learning libraries to enable FP8 support for Transformers. -As the number of parameters in Transformer models continues to grow, training and inference for -architectures such as BERT, GPT and T5 become very memory and compute-intensive. Most deep learning -frameworks train with FP32 by default. This is not essential, however, to achieve full accuracy for -many deep learning models. Using mixed-precision training, which combines single-precision (FP32) -with lower precision (e.g. FP16) format when training a model, results in significant speedups with -minimal differences in accuracy as compared to FP32 training. With Hopper GPU -architecture FP8 precision was introduced, which offers improved performance over FP16 with no -degradation in accuracy. Although all major deep learning frameworks support FP16, FP8 support is -not available natively in frameworks today. - -TE addresses the problem of FP8 support by providing APIs that integrate with popular Large Language -Model (LLM) libraries. It provides a Python API consisting of modules to easily build a Transformer -layer as well as a framework-agnostic library in C++ including structs and kernels needed for FP8 -support. Modules provided by TE internally maintain scaling factors and other values needed for FP8 -training, greatly simplifying mixed precision training for users. +As Transformer models scale to hundreds of billions of parameters across large language models, +MoE architectures, and multimodal models, training and inference become increasingly +memory and compute-intensive. Mixed-precision training, which combines single-precision (FP32) with +lower precision formats, delivers significant speedups with minimal impact on accuracy. FP8, introduced +with the Hopper GPU architecture, offers further performance gains over FP16 with no degradation in +accuracy, and newer formats like MXFP8 and NVFP4 on Blackwell push efficiency even further. + +TE integrates with popular LLM frameworks and provides optimizations that make low-precision training +work seamlessly with advanced features like MoE, tensor/sequence/context parallelism, and fused +operations. It provides a Python API consisting of modules to easily build a Transformer layer as +well as a framework-agnostic library in C++ including structs and kernels needed for FP8 support. +Modules provided by TE internally maintain scaling factors and other values needed for FP8 training, +greatly simplifying mixed precision training for users. Highlights ========== @@ -61,6 +58,7 @@ Highlights * Easy-to-use modules for building Transformer layers with FP8 support * Optimizations (e.g. fused kernels) for Transformer models * Support for FP8 on NVIDIA Hopper, Ada, and Blackwell GPUs +* Support for MXFP8 and NVFP4 on NVIDIA Blackwell GPUs * Support for optimizations across all precisions (FP16, BF16) on NVIDIA Ampere GPU architecture generations and later Examples @@ -140,7 +138,7 @@ Flax for _ in range(10): loss, (param_grads, other_grads) = fwd_bwd_fn(params, other_variables, inp) -For a more comprehensive tutorial, check out our `Getting Started Guide `_. +For a more comprehensive tutorial, check out our `Getting Started Guide `_. .. overview-end-marker-do-not-remove @@ -193,12 +191,11 @@ We recommend updating to the latest NGC container available here: * https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch * https://catalog.ngc.nvidia.com/orgs/nvidia/containers/jax -If you run any examples, please ensure you are using a matching version of TransformerEngine. TransformerEngine is pre-built and packaged inside the containers with examples available at ``/opt/transformerengine`` or ``/opt/transformer-engine``. If you would like to use examples from TE main branch and are running into import errors, please try the latest pip package or building from source, although NGC containers are recommended for ease-of-use for most users. +If you run any examples, please ensure you are using a matching version of TransformerEngine. TransformerEngine is pre-built and packaged inside the containers with examples available at ``/opt/transformerengine`` or ``/opt/transformer-engine``. **Benefits of using NGC containers:** * All dependencies pre-installed with compatible versions and optimized configurations -* NGC PyTorch 23.08+ containers include FlashAttention-2 pip Installation ^^^^^^^^^^^^^^^^ @@ -376,54 +373,43 @@ An example of this change is, False, False, True, True, True, False, False, False, False, True] -FP8 Convergence -=============== +Convergence +=========== -FP8 has been tested extensively across different model architectures and configurations and we found **no significant difference** between FP8 and BF16 training loss curves. FP8 has also been validated for accuracy on downstream LLM tasks (e.g. LAMBADA and WikiText). Below are examples of models tested for convergence across different frameworks. +FP8 and MXFP8 have been tested extensively across different model architectures and configurations and we found **no significant difference** between FP8/MXFP8 and BF16 training loss curves. FP8 and MXFP8 have also been validated for accuracy on downstream LLM tasks (e.g. LAMBADA and WikiText). Below are examples of models tested for convergence across different frameworks. +------------+------------------+---------------------------------------------------------------------------------------------------------+ | Model | Framework | Source | +============+==================+=========================================================================================================+ -| T5-770M | JAX/T5x | https://github.com/NVIDIA/JAX-Toolbox/tree/main/rosetta/rosetta/projects/t5x#convergence-and-performance| -+------------+------------------+---------------------------------------------------------------------------------------------------------+ -| MPT-1.3B | Mosaic Composer | https://www.mosaicml.com/blog/coreweave-nvidia-h100-part-1 | -+------------+------------------+---------------------------------------------------------------------------------------------------------+ -| GPT-5B | JAX/Paxml | https://github.com/NVIDIA/JAX-Toolbox/tree/main/rosetta/rosetta/projects/pax#h100-results | -+------------+------------------+---------------------------------------------------------------------------------------------------------+ -| GPT-5B | NeMo Framework | Available on request | +| MPT-1.3B | Mosaic Composer | https://www.databricks.com/blog/coreweave-nvidia-h100-part-1 | +------------+------------------+---------------------------------------------------------------------------------------------------------+ | LLama2-7B | Alibaba Pai | https://mp.weixin.qq.com/s/NQT0uKXLbXyh5031zBdeBQ | +------------+------------------+---------------------------------------------------------------------------------------------------------+ -| T5-11B | JAX/T5x | Available on request | +| LLM-8B | Megatron Core | https://arxiv.org/abs/2506.08027 | +------------+------------------+---------------------------------------------------------------------------------------------------------+ | MPT-13B | Mosaic Composer | https://www.databricks.com/blog/turbocharged-training-optimizing-databricks-mosaic-ai-stack-fp8 | +------------+------------------+---------------------------------------------------------------------------------------------------------+ -| GPT-22B | NeMo Framework | Available on request | +| MoE-16B | Megatron Core | https://arxiv.org/abs/2506.08027 | +------------+------------------+---------------------------------------------------------------------------------------------------------+ | LLama2-70B | Alibaba Pai | https://mp.weixin.qq.com/s/NQT0uKXLbXyh5031zBdeBQ | +------------+------------------+---------------------------------------------------------------------------------------------------------+ -| GPT-175B | JAX/Paxml | https://github.com/NVIDIA/JAX-Toolbox/tree/main/rosetta/rosetta/projects/pax#h100-results | -+------------+------------------+---------------------------------------------------------------------------------------------------------+ Integrations ============ Transformer Engine has been integrated with popular LLM frameworks such as: -* `DeepSpeed `_ +* `DeepSpeed `_ * `Hugging Face Accelerate `_ -* `Lightning `_ +* `Lightning `_ * `MosaicML Composer `_ * `NVIDIA JAX Toolbox `_ * `NVIDIA Megatron-LM `_ -* `NVIDIA NeMo Framework `_ +* `NVIDIA NeMo Megatron Bridge `_ * `Amazon SageMaker Model Parallel Library `_ * `Levanter `_ * `GPT-NeoX `_ -* `Hugging Face Nanotron `_ - Coming soon! -* `Colossal-AI `_ - Coming soon! -* `PeriFlow `_ - Coming soon! - +* `Hugging Face Nanotron `_ Contributing ============ @@ -442,7 +428,7 @@ Papers Videos ====== -* `Stable and Scalable FP8 Deep Learning Training on Blackwell | GTC 2025 `__ +* `Stable and Scalable FP8 Deep Learning Training on Blackwell | GTC 2025 `_ * `Blackwell Numerics for AI | GTC 2025 `_ * `Building LLMs: Accelerating Pretraining of Foundational Models With FP8 Precision | GTC 2025 `_ * `From FP8 LLM Training to Inference: Language AI at Scale | GTC 2025 `_ @@ -483,8 +469,8 @@ Previous News :alt: H200 * [11/2023] `Inflection-2: The Next Step Up `_ -* [11/2023] `Unleashing The Power Of Transformers With NVIDIA Transformer Engine `_ +* [11/2023] `Unleashing The Power Of Transformers With NVIDIA Transformer Engine `_ * [11/2023] `Accelerating PyTorch Training Workloads with FP8 `_ * [09/2023] `Transformer Engine added to AWS DL Container for PyTorch Training `_ * [06/2023] `Breaking MLPerf Training Records with NVIDIA H100 GPUs `_ -* [04/2023] `Benchmarking Large Language Models on NVIDIA H100 GPUs with CoreWeave (Part 1) `_ +* [04/2023] `Benchmarking Large Language Models on NVIDIA H100 GPUs with CoreWeave (Part 1) `_ diff --git a/UPGRADE_CONCLUSIONS_v2.17.md b/UPGRADE_CONCLUSIONS_v2.17.md new file mode 100644 index 0000000000..43cf513548 --- /dev/null +++ b/UPGRADE_CONCLUSIONS_v2.17.md @@ -0,0 +1,74 @@ +# TransformerEngine-FL v2.14 → v2.17 升级结论记录 + +本文件记录每个 Stage 完成后的结论、关键发现和决策。 + +--- + +## Stage 1: Repo Setup & Branch Preparation ✅ + +**完成时间**: 2025-01 + +**结论**: +- 三分支结构就绪 + - `base` = upstream v2.14 tag (`f031cf87`) + - `dev` = upstream v2.17 tag (`2e559f06`) + - `main` = fork 最新 (`dea7cd6c`) +- upstream remote 已添加: `https://github.com/NVIDIA/TransformerEngine.git` +- 增量规模确认: 239 commits, 482 files, +83,356/-20,066 lines + +--- + +## Stage 2: Identify Plugin Changes ✅ + +**完成时间**: 2025-01 + +### Fork 专有变更总览 + +- **base..main**: 64 commits, 233 files, +44,026/-634 lines +- Fork 变更全部是增量式（无删除/重命名 upstream 文件） + +### 文件分类 + +| 区域 | 文件数 | 类型 | 冲突风险 | +|------|--------|------|----------| +| `transformer_engine/plugin/` | ~75 | 全部新增 | 🟢 低 | +| `transformer_engine/pytorch/` | 43 修改 | patch | 🔴 高 | +| `transformer_engine/debug/` | 4 | 修改 | 🟡 中 | +| `transformer_engine/__init__.py` + `common/` | 2 | 修改 | 🟡 中 | +| `tests/` | 47 | 全部新增 | 🟢 低 | +| `.github/` | 31 | 新增+修改 | 🟢 低 | +| `qa/` | 13 | 新增+修改 | 🟡 中 | +| `build_tools/` + `setup.py` | 3 | 修改 | 🟡 中 | +| `3rdparty/` submodules | 3 | 版本更新 | 🟡 中 | +| `.gitignore` | 1 | 修改 | 🟢 低 | + +### pytorch/ 双向冲突预检 + +| 类型 | 文件数 | 处理方式 | +|------|--------|----------| +| 必然冲突（双方修改） | 38 | 手动解冲突 | +| fork-only 修改（upstream未动） | 6 | 自动保留 | +| upstream-only 修改（fork未动） | 60 | 自动接受 | + +**38个冲突文件核心模式**: fork 加了 plugin dispatch patch，upstream 改了底层实现，通常不在同一行但 git 会因上下文偏移报冲突。 + +### 3rdparty Submodule 版本决策 + +| submodule | base | fork | upstream v2.17 | 决策 | +|-----------|------|------|----------------|------| +| cudnn-frontend | `7b9b711c` | `1d6f6d9b` | `e46d7082` | 取 upstream v2.17 | +| cutlass | `57e3cfb4` | `e64a9136` | `57e3cfb4` | 保留 fork | +| googletest | `f8d7d77c` | `a0f06a70` | `f8d7d77c` | 保留 fork | +| nccl | N/A | N/A | `a6b5de08` | 接受 upstream 新增 | + +### 关键发现 + +1. Fork pytorch/ 修改本质是 plugin dispatch patch，与 upstream 实现演进通常不冲突同一行 +2. 6 个 fork-only 文件完全不会冲突 +3. upstream v2.17 新增了 `3rdparty/nccl` submodule + +--- + +## Stage 3-10 + +*待完成* diff --git a/UPGRADE_PLAN_v2.17.md b/UPGRADE_PLAN_v2.17.md new file mode 100644 index 0000000000..3a8e60b326 --- /dev/null +++ b/UPGRADE_PLAN_v2.17.md @@ -0,0 +1,653 @@ +# TransformerEngine-FL Upstream Sync Plan: v2.14 → v2.17 + +## 概览 + +| 项目 | 值 | +|------|------| +| Fork 当前基线 | upstream `v2.14` tag (commit `f031cf87`) — PR #62 已合入 main | +| 目标 | upstream `v2.17` tag (commit `2e559f06`) | +| 增量规模 | 239 commits, 482 files, +83,356 / -20,066 lines | +| csrc pybind 变化 | 34 行 `.def()` 增删（约 15-17 个 API 新增/修改） | +| main 上 v2.14 之后 fork 新增 | ~30 commits（vendor backends、CICD、FlagOS op 等） | +| 参考 PR | [#62](https://github.com/flagos-ai/TransformerEngine-FL/pull/62) (v2.9→v2.14 同步) | +| Skill 参考 | PR #67 `skills/te-fl-upstream-sync/` | + +--- + +## Stage 1: Repo Setup & Branch Preparation + +**目标**: 创建三分支结构（dev/base/main），为后续 merge 做准备。 + +**命令**: +```bash +cd /share/project/zhaoyingli/flagos/TransformerEngine-FL + +# 添加 upstream remote（如未添加） +git remote add upstream https://github.com/NVIDIA/TransformerEngine.git 2>/dev/null || true +git fetch upstream --tags + +# 创建 dev 分支 = v2.17 tag +git checkout -b dev v2.17 + +# 创建 base 分支 = v2.14 tag（PR #62 同步点） +git checkout -b base v2.14 + +# 回到 main +git checkout main +``` + +**验证**: +- `dev` → commit `2e559f06` +- `base` → commit `f031cf87` +- `main` → fork 最新 + +**产出**: SYNC_POINT.md + +--- + +## Stage 2: Identify Plugin Changes + +**目标**: 记录 fork 相对于 base 的所有专有变更，作为 Stage 3 解冲突的参考。 + +**命令**: +```bash +# Plugin 目录（全部是新增文件） +git diff --name-status base..main -- 'transformer_engine/plugin/' + +# CUDA patches / TE_DEVICE_TYPE +git diff --name-status base..main -- 'transformer_engine/__init__.py' + +# Build 系统 +git diff base..main -- setup.py CMakeLists.txt pyproject.toml + +# Python layer patches +git diff --name-status base..main -- 'transformer_engine/pytorch/' + +# 生成完整 diff 留档 +git diff base..main > /tmp/plugin_changes_full.diff +``` + +**产出**: `PLUGIN_CHANGES.md`（结构化记录新增文件、修改文件、build 变更） + +--- + +## Stage 3: Merge & Conflict Resolution + +**目标**: 将 upstream v2.17 合入 fork main，按优先级解决冲突。 + +**冲突分级**: + +| 优先级 | 范围 | 策略 | +|--------|------|------| +| P0 | `transformer_engine/pytorch/` 中 plugin 相关调用 | 保留 fork plugin dispatch, 接受 upstream 实现演进 | +| P1 | `setup.py`, `build_tools/`, `3rdparty/` | 保留 plugin 编译目标, 合入 upstream 依赖/版本更新 | +| P2 | `.github/`, `tests/`, `qa/`, `docs/`, `README` | 以 upstream 为主, 补充 fork CI 扩展 | + +**P0 解冲突原则**: +1. 读取 main 版本，保留新增的 fork 内容（plugin dispatch 调用、TE_DEVICE_TYPE 替换） +2. 接受 upstream 的功能新增和重构 +3. 检查安全/关键 bug fix + +--- + +### Stage 3.1: 创建 merge 分支并执行 merge ✅ + +```bash +git checkout main +git checkout -b merge/dev-to-main-$(date +%Y%m%d) +git merge dev --no-ff -m "merge(dev): integrate upstream v2.17" +``` + +**结果**: 分支 `merge/dev-to-main-20260807`，冲突总数 **228 个文件**: +- 173 个 fork 没改过 → Stage 3.2 自动解决 +- 55 个 fork 改过 → 手动解（Stage 3.3-3.10） + +--- + +### Stage 3.2: 自动解冲突（fork 未修改的 173 个文件） + +**分类**: + +| 分类 | 文件数 | 说明 | Stage 4 影响 | +|------|--------|------|-------------| +| Plugin 相关 | 81 | csrc pybind(12) + cpp_extensions(1) + ops/(6) + tensor/(6) + common/(56) | ⚠️ 需作为 Stage 4 API Sync 输入 | +| 非 Plugin 相关 | 92 | jax/(28) + tests/(37) + docs/(5) + build_tools/wheel(4) + qa/(5) + examples/(4) + 其他 | ✅ 无需额外处理 | + +**Plugin 相关 81 文件细分**: +- A. csrc pybind (12): pybind.cpp, extensions/*.cpp, quantizer.cpp → 直接影响 plugin ops.py +- B. cpp_extensions (1): fused_attn.py → FA Python wrapper 签名变化 +- C. ops/ (6): __init__, _common, basic/__init__, fused/__init__, backward_activation_bias, fuser +- D. tensor/ (6): _quantization_helpers, storage/*_storage.py, utils.py +- E. common/ (56): CUDA kernels (cast, fused_attn, fused_router, hadamard, gemm, normalization, multi_tensor, swizzle, transpose, triton) + +**命令**: +```bash +for file in $(git diff --name-only --diff-filter=U); do + if ! git diff base..main -- "$file" | grep -q '^[+-]'; then + git checkout --theirs "$file" && git add "$file" + echo "AUTO-RESOLVED (theirs): $file" + fi +done + +# Stage 4 准备: 导出 plugin 相关文件的 upstream diff +git diff base..dev -- transformer_engine/pytorch/csrc/ > /tmp/stage4_csrc_diff.diff +git diff base..dev -- transformer_engine/common/ > /tmp/stage4_common_diff.diff +``` + +--- + +### Stage 3.3: 解冲突 P2 — 低风险文件 + +涉及: `.github/`, `docs/`, `README.rst`, `.gitignore`, `qa/` 脚本 +策略: 以 upstream 为主，保留 fork CI 扩展 + +```bash +for file in $(git diff --name-only --diff-filter=U | grep -E '^\.(github|gitignore)|^docs/|^README|^qa/'); do + git checkout --theirs "$file" && git add "$file" + echo "P2 RESOLVED (theirs): $file" +done +``` + +--- + +### Stage 3.4: 解冲突 P1 — Build 系统 + +涉及: `setup.py`, `build_tools/pytorch.py`, `build_tools/utils.py`, `3rdparty/`, `.gitmodules` + +3rdparty submodule 决策: +- cudnn-frontend → 取 upstream v2.17 (e46d7082) +- cutlass → 保留 fork (e64a9136) +- googletest → 保留 fork (a0f06a70) +- nccl → 接受 upstream 新增 (a6b5de08) +- .gitmodules → 合并（保留原有 + 新增 nccl 条目） + +setup.py / build_tools: 保留 fork plugin 编译目标，合入 upstream 依赖更新 + +--- + +### Stage 3.5: 解冲突 P0-A — __init__.py / common / debug (6 files) + +- `transformer_engine/__init__.py` +- `transformer_engine/common/__init__.py` +- `transformer_engine/debug/features/fake_quant.py` +- `transformer_engine/debug/features/log_fp8_tensor_stats.py` +- `transformer_engine/debug/features/per_tensor_scaling.py` +- `transformer_engine/debug/features/utils/stats_buffer.py` + +策略: 保留 TE_DEVICE_TYPE 定义和 plugin import，接受 upstream 新增功能代码。 + +--- + +### Stage 3.6: 解冲突 P0-B — pytorch/attention/ (5 files) + +- `attention/dot_product_attention/backends.py` +- `attention/dot_product_attention/context_parallel.py` +- `attention/dot_product_attention/dot_product_attention.py` +- `attention/dot_product_attention/utils.py` +- `attention/multi_head_attention.py` + +策略: 保留 fork `plugin.ops.get_attention_backend()` dispatch + FlashAttentionBase，接受 upstream FA3/新参数演进。 + +--- + +### Stage 3.7: 解冲突 P0-C — pytorch/module/ (5 files) + +- `module/base.py` +- `module/grouped_linear.py` +- `module/layernorm_linear.py` +- `module/layernorm_mlp.py` +- `module/linear.py` + +策略: 保留 fork `plugin.ops.xxx()` gemm/normalization dispatch，接受 upstream quantization/precision 演进。 + +--- + +### Stage 3.8: 解冲突 P0-D — pytorch/ops/ (11 files) + +- `ops/basic/activation.py`, `basic_linear.py`, `bias.py`, `grouped_linear.py`, `swiglu.py` +- `ops/fused/forward_grouped_mlp.py`, `forward_linear_bias_activation.py`, `forward_linear_bias_add.py`, `forward_linear_scale_add.py`, `userbuffers_backward_linear.py`, `userbuffers_forward_linear.py` + +策略: 保留 fork plugin op dispatch，接受 upstream 新增 quantization path / fused op 逻辑。 + +--- + +### Stage 3.9: 解冲突 P0-E — pytorch/tensor/ (6 files) + +- `tensor/float8_blockwise_tensor.py`, `float8_tensor.py`, `grouped_tensor.py`, `mxfp8_tensor.py`, `nvfp4_tensor.py` +- `tensor/storage/grouped_tensor_storage.py` + +策略: 保留 fork TE_DEVICE_TYPE 替换，接受 upstream tensor 实现演进。 + +--- + +### Stage 3.10: 解冲突 P0-F — pytorch/ 其他文件 (11 files) + +- `pytorch/__init__.py`, `cpp_extensions/gemm.py`, `cpu_offload.py`, `distributed.py` +- `optimizers/fused_adam.py`, `permutation.py`, `quantization.py`, `setup.py` +- `transformer.py`, `triton/permutation.py`, `utils.py` + +策略: 逐文件检查 fork patch 内容，保留 plugin dispatch + TE_DEVICE_TYPE，接受 upstream 逻辑。 + +--- + +### Stage 3.11: 提交 merge commit + +```bash +# 验证无残留冲突标记 +grep -rn '<<<<<<<\|=======\|>>>>>>>' transformer_engine/ tests/ qa/ setup.py build_tools/ && echo "CONFLICT MARKERS FOUND!" || echo "CLEAN" + +# pre-commit +pre-commit run --all-files +git add -A + +# 提交 +git commit --no-edit +git log --oneline -3 +``` + +--- + +## Stage 4: Plugin API Sync + +**目标**: 让 plugin 层 1:1 覆盖 upstream csrc pybind 接口变化。 + +### 4.1 Diff csrc pybind + +```bash +git diff base..dev -- transformer_engine/pytorch/csrc/ > /tmp/csrc_diff.diff + +# 提取 ADDED APIs +grep -E '^\+.*\.def\("' /tmp/csrc_diff.diff | grep -v '^\+\+\+' | \ + sed 's/.*\.def("\([^"]*\)".*/\1/' | sort -u > /tmp/added_apis.txt + +# 提取 REMOVED APIs +grep -E '^\-.*\.def\("' /tmp/csrc_diff.diff | grep -v '^\-\-\-' | \ + sed 's/.*\.def("\([^"]*\)".*/\1/' | sort -u > /tmp/removed_apis.txt + +# 分类 +comm -23 /tmp/added_apis.txt /tmp/removed_apis.txt # 纯新增 +comm -13 /tmp/added_apis.txt /tmp/removed_apis.txt # 纯删除 +comm -12 /tmp/added_apis.txt /tmp/removed_apis.txt # 修改（签名变化） +``` + +### 4.2 更新文件列表 + +| 文件 | 操作 | +|------|------| +| `plugin/core/ops.py` | 新增/修改抽象方法 | +| `plugin/core/backends/vendor/cuda/cuda.py` | CUDA 参考实现（先改这个） | +| `plugin/core/backends/vendor/cuda/register_ops.py` | 新增 OpImpl 注册 | +| 其余 5 个 vendor（enflame/hygon/iluvatar/metax/musa） | 复制 CUDA 签名 | +| 各 vendor `register_ops.py` | 同步 OpImpl | +| `FlashAttentionBase` + 各 vendor `flash_attention.py` | 同步 forward 签名 | + +### 4.3 class-object 参数检测 + +检查 `AttentionParams` 等 dataclass 是否新增字段: +```bash +diff <(git show base:transformer_engine/pytorch/attention/dot_product_attention/utils.py | \ + sed -n '/^class AttentionParams/,/^class \|^def /p' | grep -E "^\s+\w+\s*:") \ + <(git show dev:transformer_engine/pytorch/attention/dot_product_attention/utils.py | \ + sed -n '/^class AttentionParams/,/^class \|^def /p' | grep -E "^\s+\w+\s*:") +``` + +### 4.4 Full-surface pybind 覆盖审计 + +```bash +# 提取所有 pybind export +grep -oP 'm\.def\("\K[^"]+' transformer_engine/pytorch/csrc/extensions/pybind.cpp | sort -u > /tmp/all_exports.txt + +# 对比 plugin ops.py +grep -oP 'def \K\w+' transformer_engine/plugin/core/ops.py | grep -v '^_' | sort -u > /tmp/plugin_ops.txt + +# 找缺失 +comm -23 /tmp/all_exports.txt /tmp/plugin_ops.txt +``` + +### 4.5 关键规则（PR #62 经验） + +- 所有 vendor backend **必须用显式参数签名**，禁止 `*args/**kwargs` +- enum 参数需要 `tex.DType(int(dtype))` 转换 +- quantizer 对象需要 `quantizer.dtype = tex.DType(int(qdtype))` 归一化 +- 新增 op 必须同时更新 `register_ops.py` + +--- + +## Stage 5: Patch CUDA Hardcoding + +**目标**: 将 upstream v2.14→v2.17 新引入的 `"cuda"` 硬编码替换为 `TE_DEVICE_TYPE`。 + +**扫描**: +```bash +git diff base..dev -- transformer_engine/pytorch/ ':(exclude)transformer_engine/pytorch/csrc/' \ + | grep '^+' | grep -v '^+++' \ + | grep -E 'device.*"cuda"|torch\.device\("cuda"\)|get_autocast_dtype.*"cuda"|\.device\.type.*==.*"cuda"' \ + > /tmp/cuda_string_candidates.txt +``` + +**替换规则**: + +| 模式 | 替换为 | +|------|--------| +| `device="cuda"` | `device=TE_DEVICE_TYPE` | +| `torch.device("cuda")` | `torch.device(TE_DEVICE_TYPE)` | +| `torch.get_autocast_dtype("cuda")` | `torch.get_autocast_dtype(TE_DEVICE_TYPE)` | +| `.device.type == "cuda"` | `.device.type == TE_DEVICE_TYPE` | + +**不动**: +- `torch.cuda.*` API 调用（由 vendor patches.py 运行时处理） +- `torch.cuda.CUDAGraph`、`torch.version.cuda` +- 注释/docstring 中的 `"cuda"` +- CUDA-specific guard 中的检测逻辑 + +**每个修改的文件需添加 import**: +```python +from transformer_engine import TE_DEVICE_TYPE +``` + +--- + +## Stage 6: Detect & Fix Stale References + +**目标**: 检查 fork 专有代码中引用被 upstream 重命名/移动的符号。 + +### 6.1 找出 upstream 重命名/删除 + +```bash +# 被删除的函数/类定义 +git diff base..dev -- '*.py' | grep -E '^\-(def |class )' | \ + sed 's/^-//;s/(.*//;s/def //;s/class //;s/://;s/ //g' | sort -u > /tmp/removed_symbols.txt + +# 被新增的 +git diff base..dev -- '*.py' | grep -E '^\+(def |class )' | \ + sed 's/^+//;s/(.*//;s/def //;s/class //;s/://;s/ //g' | sort -u > /tmp/added_symbols.txt + +# 真正消失的（被删且未重新添加） +comm -23 /tmp/removed_symbols.txt /tmp/added_symbols.txt > /tmp/gone_symbols.txt + +# 文件重命名/删除 +git diff --diff-filter=R --name-status -M base..dev > /tmp/renamed_files.txt +git diff --diff-filter=D --name-only base..dev > /tmp/deleted_files.txt +``` + +### 6.2 在 fork 专有代码中搜索 + +```bash +# fork 专有文件 = HEAD 有但 dev 没有的变更 +git diff --name-only dev..HEAD -- '*.py' > /tmp/fork_files.txt + +# 逐个搜索消失的符号 +while IFS= read -r symbol; do + MATCHES=$(grep -rn "$symbol" $(cat /tmp/fork_files.txt) 2>/dev/null) + [ -n "$MATCHES" ] && echo "⚠️ STALE: $symbol" && echo "$MATCHES" +done < /tmp/gone_symbols.txt +``` + +### 6.3 修复 + +- 找到新名称: `git diff base..dev -- | grep -A5 -B5 "old_name"` +- 找到新路径: `git ls-tree -r --name-only dev | grep ""` + +--- + +## Stage 7: Build & Import Verification + +**目标**: 确认合并后的代码能编译和 import。 + +```bash +git submodule update --init --recursive +pip install --no-build-isolation -e . 2>&1 | tee build.log +python -c "from transformer_engine import pytorch; print('OK')" +``` + +**预期输出**: +``` +[CUDA] Successfully loaded CUDA libs +[TE-FL manager.py INFO] OpManager initialized: N ops with M implementations +[TE-FL manager.py INFO] Registered impl_ids: ['default.flagos', 'reference.torch', 'vendor.cuda'] +``` + +**失败处理**: 根据 traceback 回溯到 Stage 3-6 修复。 + +--- + +## Stage 8: Unit & Integration Tests + +**分层执行**: + +| Level | 命令 | 验证内容 | +|-------|------|----------| +| L1 | `pytest transformer_engine/plugin/tests/ -k plugin -v` | Plugin 注册/dispatch | +| L2 | `pytest tests/pytorch/ -v` | 上游集成测试 | +| L2.5a | `TE_PATH=$(pwd) bash qa/L0_pytorch_debug_unittest/test.sh` | Debug 单测 | +| L2.5b | `TE_PATH=$(pwd) bash qa/L0_pytorch_unittest/test.sh` | PyTorch 单测 | +| L2.5c | `TE_PATH=$(pwd) bash qa/L1_pytorch_distributed_unittest/test.sh` | 分布式测试 | +| L2.6 | `python transformer_engine/plugin/tests/run_all_tests.py` | Plugin 全量测试 | +| L3 | `python tests/pytorch/test_sanity.py` | E2E sanity | + +**前置检查**: CI 脚本引用的测试文件是否因 upstream rename 而失效: +```bash +grep -oP '(?<=\$TE_PATH/)tests/pytorch/[^\s"]+\.py' qa/L0_pytorch_unittest/test.sh | \ + sort -u | while read f; do [ ! -f "$f" ] && echo "MISSING: $f"; done +``` + +**PR #62 踩过的坑**: +- `test_float8tensor.py` → `test_quantized_tensor.py`（upstream rename） +- OP API 签名不匹配（如 `fused_topk_with_score_function_bwd` 缺参数） +- MetaX 平台特有跳过项需维护 + +--- + +## Stage 5: Patch CUDA Hardcoding + +**目标**: 将 upstream v2.14→v2.17 新引入的 `"cuda"` 硬编码替换为 `TE_DEVICE_TYPE`。 + +**扫描**: +```bash +git diff base..dev -- transformer_engine/pytorch/ ':(exclude)transformer_engine/pytorch/csrc/' \ + | grep '^+' | grep -v '^+++' \ + | grep -E 'device.*"cuda"|torch\.device\("cuda"\)|get_autocast_dtype.*"cuda"|\.device\.type.*==.*"cuda"' \ + > /tmp/cuda_string_candidates.txt +``` + +**替换规则**: + +| 模式 | 替换为 | +|------|--------| +| `device="cuda"` | `device=TE_DEVICE_TYPE` | +| `torch.device("cuda")` | `torch.device(TE_DEVICE_TYPE)` | +| `torch.get_autocast_dtype("cuda")` | `torch.get_autocast_dtype(TE_DEVICE_TYPE)` | +| `.device.type == "cuda"` | `.device.type == TE_DEVICE_TYPE` | + +**不动**: +- `torch.cuda.*` API 调用（由 vendor patches.py 运行时处理） +- `torch.cuda.CUDAGraph`、`torch.version.cuda` +- 注释/docstring 中的 `"cuda"` +- CUDA-specific guard 中的检测逻辑 + +**每个修改的文件需添加 import**: +```python +from transformer_engine import TE_DEVICE_TYPE +``` + +--- + +## Stage 6: Detect & Fix Stale References + +**目标**: 检查 fork 专有代码中引用被 upstream 重命名/移动的符号。 + +### 6.1 找出 upstream 重命名/删除 + +```bash +# 被删除的函数/类定义 +git diff base..dev -- '*.py' | grep -E '^\-(def |class )' | \ + sed 's/^-//;s/(.*//;s/def //;s/class //;s/://;s/ //g' | sort -u > /tmp/removed_symbols.txt + +# 被新增的 +git diff base..dev -- '*.py' | grep -E '^\+(def |class )' | \ + sed 's/^+//;s/(.*//;s/def //;s/class //;s/://;s/ //g' | sort -u > /tmp/added_symbols.txt + +# 真正消失的（被删且未重新添加） +comm -23 /tmp/removed_symbols.txt /tmp/added_symbols.txt > /tmp/gone_symbols.txt + +# 文件重命名/删除 +git diff --diff-filter=R --name-status -M base..dev > /tmp/renamed_files.txt +git diff --diff-filter=D --name-only base..dev > /tmp/deleted_files.txt +``` + +### 6.2 在 fork 专有代码中搜索 + +```bash +# fork 专有文件 = HEAD 有但 dev 没有的变更 +git diff --name-only dev..HEAD -- '*.py' > /tmp/fork_files.txt + +# 逐个搜索消失的符号 +while IFS= read -r symbol; do + MATCHES=$(grep -rn "$symbol" $(cat /tmp/fork_files.txt) 2>/dev/null) + [ -n "$MATCHES" ] && echo "⚠️ STALE: $symbol" && echo "$MATCHES" +done < /tmp/gone_symbols.txt +``` + +### 6.3 修复 + +- 找到新名称: `git diff base..dev -- | grep -A5 -B5 "old_name"` +- 找到新路径: `git ls-tree -r --name-only dev | grep ""` + +--- + +## Stage 7: Build & Import Verification + +**目标**: 确认合并后的代码能编译和 import。 + +```bash +git submodule update --init --recursive +pip install --no-build-isolation -e . 2>&1 | tee build.log +python -c "from transformer_engine import pytorch; print('OK')" +``` + +**预期输出**: +``` +[CUDA] Successfully loaded CUDA libs +[TE-FL manager.py INFO] OpManager initialized: N ops with M implementations +[TE-FL manager.py INFO] Registered impl_ids: ['default.flagos', 'reference.torch', 'vendor.cuda'] +``` + +**失败处理**: 根据 traceback 回溯到 Stage 3-6 修复。 + +--- + +## Stage 8: Unit & Integration Tests + +**分层执行**: + +| Level | 命令 | 验证内容 | +|-------|------|----------| +| L1 | `pytest transformer_engine/plugin/tests/ -k plugin -v` | Plugin 注册/dispatch | +| L2 | `pytest tests/pytorch/ -v` | 上游集成测试 | +| L2.5a | `TE_PATH=$(pwd) bash qa/L0_pytorch_debug_unittest/test.sh` | Debug 单测 | +| L2.5b | `TE_PATH=$(pwd) bash qa/L0_pytorch_unittest/test.sh` | PyTorch 单测 | +| L2.5c | `TE_PATH=$(pwd) bash qa/L1_pytorch_distributed_unittest/test.sh` | 分布式测试 | +| L2.6 | `python transformer_engine/plugin/tests/run_all_tests.py` | Plugin 全量测试 | +| L3 | `python tests/pytorch/test_sanity.py` | E2E sanity | + +**前置检查**: CI 脚本引用的测试文件是否因 upstream rename 而失效: +```bash +grep -oP '(?<=\$TE_PATH/)tests/pytorch/[^\s"]+\.py' qa/L0_pytorch_unittest/test.sh | \ + sort -u | while read f; do [ ! -f "$f" ] && echo "MISSING: $f"; done +``` + +**PR #62 踩过的坑**: +- `test_float8tensor.py` → `test_quantized_tensor.py`（upstream rename） +- OP API 签名不匹配（如 `fused_topk_with_score_function_bwd` 缺参数） +- MetaX 平台特有跳过项需维护 + +--- + +## Stage 9: Merge to Main (Tree Replacement) + +**目标**: 将完成验证的 merge 分支通过 tree replacement 策略合入 main，提交 PR。 + +**为什么用 tree replacement 而不是普通 merge**: +`-X theirs` 只解决冲突，非冲突变更仍会合并两侧。当 fork 和 upstream 独立添加了相同 patch 时， +git 会保留两份副本导致重复代码。Tree replacement 完全避免此问题。 + +**命令**: +```bash +git checkout main && git pull origin main +git checkout -b merge-to-main-$(date +%Y%m%d) + +# Tree replacement: 记录两个 parent 但用 merge 分支的 tree +git merge -s ours merge/dev-to-main-YYYYMMDD --no-edit +git read-tree -m -u merge/dev-to-main-YYYYMMDD + +# pre-commit 格式化 +pre-commit run --all-files +git add -A && git commit --amend --no-edit +``` + +**清理中间文件**: +```bash +for f in SYNC_POINT.md MERGE_RECORD.md UPSTREAM_SYNC.md; do + [ -f "$f" ] && git rm "$f" +done +git diff --cached --quiet || git commit -m "chore: remove intermediate sync record files" +``` + +**验证**: +```bash +git diff merge/dev-to-main-YYYYMMDD HEAD --stat +git log --oneline --graph -5 +pip install --no-build-isolation -e . +python -c "import transformer_engine; print('OK')" +``` + +--- + +## Stage 10: FlagScale E2E Training Validation + +**目标**: 在 FlagScale 真实训练场景下验证合并后的 TE-FL。 + +**测试矩阵**（复用 PR #62 标准）: + +### Qwen3-32B（16 layers, 20 iters, 1 node × 8 GPUs） + +| Config | 预期 | +|--------|------| +| vendor-flash / vendor-fused / vendor-unfused | PASS | +| flagos-flash / flagos-unfused | PASS | +| reference-flash / reference-unfused | PASS | + +### DeepSeek-V3 16BA3B（18 layers + 1 mtp, 20 iters） + +| Config | 预期 | +|--------|------| +| vendor-unfused / flagos-unfused / reference-unfused | PASS | + +**成功标准**: 至少一个组合跑完 20 步无错误，loss 下降。 + +--- + +## 风险与注意事项 + +| 风险 | 应对 | +|------|------| +| NVFP4/MXFP8 新 op 多 | Stage 4 重点关注 | +| FA3 参数变化 | Stage 4 FlashAttentionBase 同步 | +| 新 vendor stale refs | Stage 6 重点扫描 | +| Build system 演进 | Stage 3 P1 仔细处理 | + +## 时间估算 + +| 阶段 | 预计耗时 | +|------|----------| +| Stage 1-2 | 0.5 天 | +| Stage 3 (merge) | 2-3 天 | +| Stage 4 (Plugin API) | 2-3 天 | +| Stage 5-6 | 1 天 | +| Stage 7-8 (build+tests) | 1-2 天 | +| Stage 9-10 | 1 天 | +| **总计** | **7-10 天** | + +## Rollback + +任何时候出问题: `git revert -m 1 ` diff --git a/benchmarks/benchmark_rht_cast.py b/benchmarks/benchmark_rht_cast.py index badab1d199..2ef63f705a 100644 --- a/benchmarks/benchmark_rht_cast.py +++ b/benchmarks/benchmark_rht_cast.py @@ -8,7 +8,6 @@ import torch.utils.benchmark as benchmark import transformer_engine.pytorch as te -import transformer_engine_torch as tex import transformer_engine.pytorch.cpp_extensions as ext from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer @@ -17,7 +16,7 @@ permute_scale = False TORCH_TO_TE_FLOAT_MAP = { - torch.bfloat16: tex.DType.kBFloat16, + torch.bfloat16: te.DType.kBFloat16, } @@ -31,7 +30,7 @@ def run_kernel(shape, stochastic_rounding: bool, input_dtype=torch.bfloat16): # Quantize nvfp4_quantizer = NVFP4Quantizer( - fp4_dtype=tex.DType.kFloat4E2M1, + fp4_dtype=te.DType.kFloat4E2M1, rowwise=True, columnwise=True, with_amax_reduction=False, diff --git a/benchmarks/benchmark_rht_cast_swizzle_fusion.py b/benchmarks/benchmark_rht_cast_swizzle_fusion.py new file mode 100644 index 0000000000..13f264862c --- /dev/null +++ b/benchmarks/benchmark_rht_cast_swizzle_fusion.py @@ -0,0 +1,189 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Benchmark NVFP4 RHT cast-fusion with vs without fused GEMM-swizzled SF output. + +For each shape we measure two paths and two builds: + + * path = "quant_only": just NVFP4Quantizer(x) + * path = "quant_plus_swizzle": NVFP4Quantizer(x) + tex.swizzle_scales_for_gemm_(t) + (this is what te.Linear -> tex.generic_gemm does right before the + cuBLAS LT NVFP4 GEMM dispatch) + + * build = "baseline": optimize_for_gemm=False + -> quant kernel emits compact SF; + tex.swizzle_scales_for_gemm_ launches the standalone + swizzle_{row,col}_scaling_kernel pass before GEMM. + * build = "swizzle_fusion": optimize_for_gemm=True + -> quant kernel emits GEMM-swizzled SF directly (via the + kEnableSwizzleSFOutput compile-time switch in + row_cast_col_hadamard_transform_cast_fusion.cu); + tex.swizzle_scales_for_gemm_ early-returns and the standalone + swizzle pass disappears from the timeline. + +The wall-clock delta on the "quant_plus_swizzle" path is the production +saving of this PR. +""" + +import argparse +import torch +import pandas as pd +import torch.utils.benchmark as benchmark + +import transformer_engine.pytorch as te # noqa: F401 must be first per te-python-import-order +import transformer_engine_torch as tex +from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + + +def make_quantizer(optimize_for_gemm: bool) -> NVFP4Quantizer: + q = NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=True, + with_post_rht_amax=True, + with_random_sign_mask=True, + ) + q.optimize_for_gemm = optimize_for_gemm + return q + + +def _bench(stmt: str, globals_dict: dict, min_run_time: float) -> float: + """Returns median wall-clock per call in microseconds.""" + timing = benchmark.Timer( + stmt=stmt, + globals=globals_dict, + num_threads=1, + ).blocked_autorange(min_run_time=min_run_time) + return timing.median * 1e6 + + +def run_shape(shape, min_run_time: float): + M, K = shape + assert M % 16 == 0 and K % 16 == 0, "Shape must be divisible by 16" + + x = torch.randn([M, K], dtype=torch.bfloat16, device="cuda") + q_base = make_quantizer(optimize_for_gemm=False) + q_swf = make_quantizer(optimize_for_gemm=True) + + # quant_only path + quant_only_base_us = _bench( + stmt="q(x)", + globals_dict={"q": q_base, "x": x}, + min_run_time=min_run_time, + ) + quant_only_swf_us = _bench( + stmt="q(x)", + globals_dict={"q": q_swf, "x": x}, + min_run_time=min_run_time, + ) + + # quant_plus_swizzle path (this is what te.Linear actually runs) + quant_plus_swizzle_base_us = _bench( + stmt="t = q(x); tex.swizzle_scales_for_gemm_(t)", + globals_dict={"q": q_base, "x": x, "tex": tex}, + min_run_time=min_run_time, + ) + quant_plus_swizzle_swf_us = _bench( + stmt="t = q(x); tex.swizzle_scales_for_gemm_(t)", + globals_dict={"q": q_swf, "x": x, "tex": tex}, + min_run_time=min_run_time, + ) + + saved_us = quant_plus_swizzle_base_us - quant_plus_swizzle_swf_us + speedup = ( + quant_plus_swizzle_base_us / quant_plus_swizzle_swf_us + if quant_plus_swizzle_swf_us > 0 + else float("inf") + ) + + print( + f" shape={shape}: quant_only base={quant_only_base_us:.2f}us, " + f"SUT={quant_only_swf_us:.2f}us | " + f"quant+swizzle base={quant_plus_swizzle_base_us:.2f}us, " + f"SUT={quant_plus_swizzle_swf_us:.2f}us " + f"-> saved {saved_us:.2f}us ({speedup:.2f}x)" + ) + + return { + "shape": shape, + "M": M, + "K": K, + "quant_only_base_us": quant_only_base_us, + "quant_only_swf_us": quant_only_swf_us, + "quant_plus_swizzle_base_us": quant_plus_swizzle_base_us, + "quant_plus_swizzle_swf_us": quant_plus_swizzle_swf_us, + "saved_us": saved_us, + "speedup": speedup, + } + + +# Nsight Compute Profiling Command (for verifying the swizzle kernel disappears): +# ncu -f -o swizzle_fusion --set=full \ +# --kernel-name "regex:swizzle_(row|col)_scaling_kernel|cast_col_hadamard_transform_cast_fusion" \ +# -s 5 -c 10 python benchmarks/benchmark_rht_cast_swizzle_fusion.py --profile + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--profile", + action="store_true", + help="Run only one shape for use with ncu/nsys; longer min_run_time", + ) + parser.add_argument( + "--min-run-time", + type=float, + default=2.0, + help="Minimum total measured time per cell in seconds (benchmark.Timer)", + ) + parser.add_argument( + "--csv", + type=str, + default="benchmark_rht_cast_swizzle_fusion.csv", + help="CSV output path", + ) + args = parser.parse_args() + + if args.profile: + print("Profiling mode enabled (single shape).") + shapes = [(8192, 4096)] + min_run_time = max(5.0, args.min_run_time) + else: + shapes = [ + # production-class shapes + (8192, 5120), + (8192, 10240), + (8192, 2560), + (8192, 11328), + (8192, 3584), + (5120, 8192), + (10240, 8192), + (2560, 8192), + (11328, 8192), + (3584, 8192), + (4096, 16384), + (14336, 16384), + ] + min_run_time = args.min_run_time + + print( + "NVFP4 RHT cast-fusion: swizzle-fusion (optimize_for_gemm=True) vs baseline. " + f"min_run_time={min_run_time}s per cell, BF16 input, " + "rowwise+columnwise SF, RHT=True+post_rht_amax." + ) + rows = [] + for shape in shapes: + print(f"Running {shape} ...") + rows.append(run_shape(shape, min_run_time)) + + df = pd.DataFrame(rows) + pd.set_option("display.max_columns", None) + pd.set_option("display.width", 200) + print() + print(df.to_string(index=False)) + df.to_csv(args.csv, index=False) + print(f"\nWrote {args.csv}") diff --git a/benchmarks/gemm/benchmark_gemm.py b/benchmarks/gemm/benchmark_gemm.py new file mode 100644 index 0000000000..2382cc339f --- /dev/null +++ b/benchmarks/gemm/benchmark_gemm.py @@ -0,0 +1,1883 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + + +"""Unified GEMM benchmark for BF16, FP8 (Current/Delayed/Block), MXFP8, and NVFP4 precisions. + +Compares matrix-multiplication throughput across precisions using +Transformer Engine on NVIDIA GPUs. Supports two timing back-ends, +pre-quantized and autocast quantization modes, arbitrary MxKxN matrix +shapes, Nsight Systems profiling integration, and bar-chart output. + +Timing back-ends +---------------- +* **cuda-events** -- CUDA event pairs with a leading-kernel trick to + hide CPU dispatch latency. Measures the full GPU-side duration of + the timed loop (includes quantisation when using autocast mode). +* **profiler** -- ``torch.profiler`` (CUPTI) kernel timestamps. + Only the matched GEMM compute kernels (gemm, nvjet, xmma, cutlass) + are summed, giving a kernel-only measurement. + +Usage examples:: + + # Kernel-only timing via torch.profiler: + python benchmarks/gemm/benchmark_gemm.py --timing profiler --pre-quantize -o kernel.png + + # End-to-end timing via CUDA events: + python benchmarks/gemm/benchmark_gemm.py --timing cuda-events -o e2e.png + + # Custom non-square shapes: + python benchmarks/gemm/benchmark_gemm.py --shapes 88064x2560x10240,88064x10240x2560 + + # Nsight profiling of a single shape: + nsys profile --capture-range=cudaProfilerApi \\ + python benchmarks/gemm/benchmark_gemm.py --profile --profile-shape 4096 + + # Model config mode (derives all 12 GEMM shapes from hyperparameters): + python benchmarks/gemm/benchmark_gemm.py \\ + --hidden_size 4096 --intermediate_size 16384 \\ + --num_attention_heads 32 --num_hidden_layers 24 \\ + --micro_batch_size 31 --sequence_length 512 +""" + +import argparse +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +import matplotlib.pyplot as plt +import numpy as np +import torch +from torch.profiler import ProfilerActivity, profile + +try: + import transformer_engine.pytorch as te + import transformer_engine_torch as tex + from transformer_engine.common.recipe import ( + DelayedScaling, + Float8BlockScaling, + Float8CurrentScaling, + Format, + MXFP8BlockScaling, + NVFP4BlockScaling, + ) + + TE_AVAILABLE = True +except ImportError: + TE_AVAILABLE = False + + +GEMM_KERNEL_PATTERNS = ("gemm", "nvjet", "xmma", "cutlass") + +PRECISION_COLORS = { + "BF16": "#808080", + "FP8Current": "#2E8B57", + "FP8Delayed": "#20B2AA", + "FP8Block": "#006400", + "MXFP8": "#4B0082", + "NVFP4": "#B22222", +} + + +# --------------------------------------------------------------------------- +# Data structures +# --------------------------------------------------------------------------- +@dataclass +class GEMMResult: + """Single GEMM benchmark measurement.""" + + tflops: float + avg_time_ms: float + shape: tuple[int, int, int] + precision: str + + +@dataclass +class ModelConfig: + """Transformer model hyperparameters for GEMM shape derivation.""" + + hidden_size: int + intermediate_size: int + num_attention_heads: int + num_hidden_layers: int + micro_batch_size: int + sequence_length: int + + +# --------------------------------------------------------------------------- +# Hardware helpers +# --------------------------------------------------------------------------- +def is_blackwell_available() -> bool: + """Return True when the current device is Blackwell (SM100+) for NVFP4 support.""" + if not torch.cuda.is_available(): + return False + major, _ = torch.cuda.get_device_capability() + return major >= 10 + + +def compute_gemm_flops(M: int, K: int, N: int) -> int: + """Theoretical FLOP count for C = A @ B: 2 * M * N * K.""" + return 2 * M * N * K + + +# --------------------------------------------------------------------------- +# torch.profiler helpers (kernel-only timing) +# --------------------------------------------------------------------------- +def _is_gemm_kernel(name: str) -> bool: + """Return True when *name* looks like a GEMM compute kernel.""" + low = name.lower() + return any(p in low for p in GEMM_KERNEL_PATTERNS) + + +def _extract_gemm_kernel_time_us( + prof_result: profile, + num_iters: int, + verbose: bool = False, +) -> float: + """Average GEMM-kernel time in microseconds from profiler events.""" + total_us = 0.0 + count = 0 + seen: dict[str, float] = {} + + for evt in prof_result.events(): + if evt.device_type == torch.autograd.DeviceType.CUDA and _is_gemm_kernel(evt.name): + total_us += evt.device_time + count += 1 + seen[evt.name] = seen.get(evt.name, 0.0) + evt.device_time + + if verbose and seen: + print(f" Matched GEMM kernels ({count} invocations):") + for kname, kus in seen.items(): + print(f" {kname}: {kus:.0f} us total") + + if count == 0: + if verbose: + print(" WARNING: No GEMM kernels found. All CUDA events:") + for evt in prof_result.events(): + if evt.device_type == torch.autograd.DeviceType.CUDA: + print(f" {evt.name}: {evt.device_time:.0f} us") + return 0.0 + + return total_us / num_iters + + +# --------------------------------------------------------------------------- +# Timing wrappers +# --------------------------------------------------------------------------- +def _time_with_profiler( + run_fn, + num_iters: int, + flops: int, + verbose: bool = False, +) -> tuple[float, float]: + """Return (tflops, avg_ms) using torch.profiler kernel extraction.""" + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for _ in range(num_iters): + run_fn() + torch.cuda.synchronize() + + avg_us = _extract_gemm_kernel_time_us(prof, num_iters, verbose=verbose) + avg_s = avg_us / 1e6 + tflops = (flops / avg_s) / 1e12 if avg_s > 0 else 0.0 + return tflops, avg_us / 1000.0 + + +def _time_with_cuda_events( + run_fn, + num_iters: int, + flops: int, + leading_fn=None, +) -> tuple[float, float]: + """Return (tflops, avg_ms) using CUDA events with optional leading kernel.""" + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + if leading_fn is not None: + leading_fn() + + start.record() + for _ in range(num_iters): + run_fn() + end.record() + torch.cuda.synchronize() + + avg_ms = start.elapsed_time(end) / num_iters + avg_s = avg_ms / 1000.0 + tflops = (flops / avg_s) / 1e12 if avg_s > 0 else 0.0 + return tflops, avg_ms + + +# --------------------------------------------------------------------------- +# BF16 benchmark +# --------------------------------------------------------------------------- +def benchmark_bf16( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> GEMMResult: + """Benchmark BF16 torch.matmul.""" + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + A = torch.randn(M, K, dtype=torch.bfloat16, device=device) + B = torch.randn(K, N, dtype=torch.bfloat16, device=device) + + for _ in range(num_warmup): + torch.matmul(A, B) + torch.cuda.synchronize() + + def _run(): + torch.matmul(A, B) + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + A_lg = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + B_lg = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + tflops, avg_ms = _time_with_cuda_events( + _run, num_iters, flops, leading_fn=lambda: torch.matmul(A_lg, B_lg) + ) + del A_lg, B_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="BF16") + + +# --------------------------------------------------------------------------- +# FP8 tensor-wise scaling benchmarks (CurrentScaling / DelayedScaling) +# --------------------------------------------------------------------------- +def benchmark_fp8_current( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """FP8 GEMM with Float8CurrentScaling recipe via te.Linear autocast.""" + if not TE_AVAILABLE: + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + linear = te.Linear(K, N, bias=False, params_dtype=torch.bfloat16).to(device) + x = torch.randn(M, K, dtype=torch.bfloat16, device=device) + recipe = Float8CurrentScaling() + + with te.autocast(enabled=True, recipe=recipe): + for _ in range(num_warmup): + linear(x) + torch.cuda.synchronize() + + def _run(): + linear(x) + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + lin_lg = te.Linear(4096, 4096, bias=False, params_dtype=torch.bfloat16).to(device) + x_lg = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + tflops, avg_ms = _time_with_cuda_events( + _run, num_iters, flops, leading_fn=lambda: lin_lg(x_lg) + ) + del lin_lg, x_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="FP8Current") + + +def benchmark_fp8_current_prequantized( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """Pre-quantized FP8 GEMM with Float8CurrentScaling via tex.generic_gemm.""" + if not TE_AVAILABLE: + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + try: + quantizer = te.Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device=device) + + A_q = quantizer.quantize(torch.randn(K, M, dtype=torch.bfloat16, device=device)) + B_q = quantizer.quantize(torch.randn(K, N, dtype=torch.bfloat16, device=device)) + D = torch.empty(N, M, dtype=torch.bfloat16, device=device) + ws_size = 32 * 1024 * 1024 + ws = torch.empty(ws_size, dtype=torch.uint8, device=device) + + def _run(): + tex.generic_gemm( + A_q, + False, + B_q, + True, + D, + None, + tex.DType.kBFloat16, + None, + tex.DType.kBFloat16, + False, + None, + False, + ws, + ws_size, + False, + False, + ) + + for _ in range(num_warmup): + _run() + torch.cuda.synchronize() + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + A_lg_q = quantizer.quantize( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + ) + B_lg_q = quantizer.quantize( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + ) + D_lg = torch.empty(4096, 4096, dtype=torch.bfloat16, device=device) + + def _lead(): + tex.generic_gemm( + A_lg_q, + False, + B_lg_q, + True, + D_lg, + None, + tex.DType.kBFloat16, + None, + tex.DType.kBFloat16, + False, + None, + False, + ws, + ws_size, + False, + False, + ) + + tflops, avg_ms = _time_with_cuda_events(_run, num_iters, flops, leading_fn=_lead) + del A_lg_q, B_lg_q, D_lg + + return GEMMResult( + tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="FP8Current" + ) + except Exception as e: + print(f"Warning: FP8 CurrentScaling prequantized benchmark failed: {e}") + return None + + +def benchmark_fp8_delayed( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """FP8 GEMM with DelayedScaling recipe via te.Linear autocast.""" + if not TE_AVAILABLE: + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + linear = te.Linear(K, N, bias=False, params_dtype=torch.bfloat16).to(device) + x = torch.randn(M, K, dtype=torch.bfloat16, device=device) + recipe = DelayedScaling() + + with te.autocast(enabled=True, recipe=recipe): + for _ in range(num_warmup): + linear(x) + torch.cuda.synchronize() + + def _run(): + linear(x) + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + lin_lg = te.Linear(4096, 4096, bias=False, params_dtype=torch.bfloat16).to(device) + x_lg = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + tflops, avg_ms = _time_with_cuda_events( + _run, num_iters, flops, leading_fn=lambda: lin_lg(x_lg) + ) + del lin_lg, x_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="FP8Delayed") + + +# --------------------------------------------------------------------------- +# MXFP8 benchmarks +# --------------------------------------------------------------------------- +def benchmark_fp8( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """MXFP8 GEMM via te.Linear autocast.""" + if not TE_AVAILABLE: + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + linear = te.Linear(K, N, bias=False, params_dtype=torch.bfloat16).to(device) + x = torch.randn(M, K, dtype=torch.bfloat16, device=device) + recipe = MXFP8BlockScaling(fp8_format=Format.E4M3) + + with te.autocast(enabled=True, recipe=recipe): + for _ in range(num_warmup): + linear(x) + torch.cuda.synchronize() + + def _run(): + linear(x) + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + lin_lg = te.Linear(4096, 4096, bias=False, params_dtype=torch.bfloat16).to(device) + x_lg = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + tflops, avg_ms = _time_with_cuda_events( + _run, num_iters, flops, leading_fn=lambda: lin_lg(x_lg) + ) + del lin_lg, x_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="MXFP8") + + +def benchmark_fp8_prequantized( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """Pre-quantized MXFP8 GEMM via tex.generic_gemm (raw kernel throughput).""" + if not TE_AVAILABLE: + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + try: + quantizer = te.MXFP8Quantizer(tex.DType.kFloat8E4M3) + + # tex.generic_gemm uses column-major convention: A=(K,M), B=(K,N), + # D=(N,M) with transa=False, transb=True for a logical C(M,N) GEMM. + A_q = quantizer.quantize(torch.randn(K, M, dtype=torch.bfloat16, device=device)) + B_q = quantizer.quantize(torch.randn(K, N, dtype=torch.bfloat16, device=device)) + D = torch.empty(N, M, dtype=torch.bfloat16, device=device) + ws_size = 32 * 1024 * 1024 + ws = torch.empty(ws_size, dtype=torch.uint8, device=device) + + def _run(): + tex.generic_gemm( + A_q, + False, + B_q, + True, + D, + None, + tex.DType.kBFloat16, + None, + tex.DType.kBFloat16, + False, + None, + False, + ws, + ws_size, + False, + False, + ) + + for _ in range(num_warmup): + _run() + torch.cuda.synchronize() + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + A_lg_q = quantizer.quantize( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + ) + B_lg_q = quantizer.quantize( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + ) + D_lg = torch.empty(4096, 4096, dtype=torch.bfloat16, device=device) + + def _lead(): + tex.generic_gemm( + A_lg_q, + False, + B_lg_q, + True, + D_lg, + None, + tex.DType.kBFloat16, + None, + tex.DType.kBFloat16, + False, + None, + False, + ws, + ws_size, + False, + False, + ) + + tflops, avg_ms = _time_with_cuda_events(_run, num_iters, flops, leading_fn=_lead) + del A_lg_q, B_lg_q, D_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="MXFP8") + except Exception as e: + print(f"Warning: FP8 prequantized benchmark failed: {e}") + return None + + +# --------------------------------------------------------------------------- +# Float8 Block-Scaling benchmarks +# --------------------------------------------------------------------------- +def benchmark_fp8_block( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """FP8 GEMM with Float8BlockScaling recipe via te.Linear autocast.""" + if not TE_AVAILABLE: + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + linear = te.Linear(K, N, bias=False, params_dtype=torch.bfloat16).to(device) + x = torch.randn(M, K, dtype=torch.bfloat16, device=device) + recipe = Float8BlockScaling(fp8_format=Format.E4M3) + + with te.autocast(enabled=True, recipe=recipe): + for _ in range(num_warmup): + linear(x) + torch.cuda.synchronize() + + def _run(): + linear(x) + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + lin_lg = te.Linear(4096, 4096, bias=False, params_dtype=torch.bfloat16).to(device) + x_lg = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + tflops, avg_ms = _time_with_cuda_events( + _run, num_iters, flops, leading_fn=lambda: lin_lg(x_lg) + ) + del lin_lg, x_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="FP8Block") + + +def benchmark_fp8_block_prequantized( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """Pre-quantized FP8 GEMM with Float8BlockScaling via tex.generic_gemm.""" + if not TE_AVAILABLE: + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + try: + quantizer = te.Float8BlockQuantizer( + tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + ) + + A_q = quantizer.quantize(torch.randn(K, M, dtype=torch.bfloat16, device=device)) + B_q = quantizer.quantize(torch.randn(K, N, dtype=torch.bfloat16, device=device)) + D = torch.empty(N, M, dtype=torch.bfloat16, device=device) + ws_size = 32 * 1024 * 1024 + ws = torch.empty(ws_size, dtype=torch.uint8, device=device) + + def _run(): + tex.generic_gemm( + A_q, + False, + B_q, + True, + D, + None, + tex.DType.kBFloat16, + None, + tex.DType.kBFloat16, + False, + None, + False, + ws, + ws_size, + False, + False, + ) + + for _ in range(num_warmup): + _run() + torch.cuda.synchronize() + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + A_lg_q = quantizer.quantize( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + ) + B_lg_q = quantizer.quantize( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + ) + D_lg = torch.empty(4096, 4096, dtype=torch.bfloat16, device=device) + + def _lead(): + tex.generic_gemm( + A_lg_q, + False, + B_lg_q, + True, + D_lg, + None, + tex.DType.kBFloat16, + None, + tex.DType.kBFloat16, + False, + None, + False, + ws, + ws_size, + False, + False, + ) + + tflops, avg_ms = _time_with_cuda_events(_run, num_iters, flops, leading_fn=_lead) + del A_lg_q, B_lg_q, D_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="FP8Block") + except Exception as e: + print(f"Warning: FP8 Block-Scaling prequantized benchmark failed: {e}") + return None + + +# --------------------------------------------------------------------------- +# NVFP4 benchmarks (Blackwell SM100+ only) +# --------------------------------------------------------------------------- +def benchmark_fp4( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """NVFP4 GEMM via te.Linear autocast (Blackwell only).""" + if not TE_AVAILABLE or not is_blackwell_available(): + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + linear = te.Linear(K, N, bias=False, params_dtype=torch.bfloat16).to(device) + x = torch.randn(M, K, dtype=torch.bfloat16, device=device) + recipe = NVFP4BlockScaling(fp4_format=Format.E2M1) + + with te.autocast(enabled=True, recipe=recipe): + for _ in range(num_warmup): + linear(x) + torch.cuda.synchronize() + + def _run(): + linear(x) + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + lin_lg = te.Linear(4096, 4096, bias=False, params_dtype=torch.bfloat16).to(device) + x_lg = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + tflops, avg_ms = _time_with_cuda_events( + _run, num_iters, flops, leading_fn=lambda: lin_lg(x_lg) + ) + del lin_lg, x_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="NVFP4") + + +def benchmark_fp4_prequantized( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """Pre-quantized NVFP4 GEMM via tex.generic_gemm (Blackwell only).""" + if not TE_AVAILABLE or not is_blackwell_available(): + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + try: + quantizer = te.NVFP4Quantizer(tex.DType.kFloat4E2M1) + + # tex.generic_gemm uses column-major convention: A=(K,M), B=(K,N), + # D=(N,M) with transa=False, transb=True for a logical C(M,N) GEMM. + A_q = quantizer.quantize(torch.randn(K, M, dtype=torch.bfloat16, device=device)) + B_q = quantizer.quantize(torch.randn(K, N, dtype=torch.bfloat16, device=device)) + D = torch.empty(N, M, dtype=torch.bfloat16, device=device) + ws_size = 32 * 1024 * 1024 + ws = torch.empty(ws_size, dtype=torch.uint8, device=device) + + def _run(): + tex.generic_gemm( + A_q, + False, + B_q, + True, + D, + None, + tex.DType.kBFloat16, + None, + tex.DType.kBFloat16, + False, + None, + False, + ws, + ws_size, + False, + False, + ) + + for _ in range(num_warmup): + _run() + torch.cuda.synchronize() + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + A_lg_q = quantizer.quantize( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + ) + B_lg_q = quantizer.quantize( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + ) + D_lg = torch.empty(4096, 4096, dtype=torch.bfloat16, device=device) + + def _lead(): + tex.generic_gemm( + A_lg_q, + False, + B_lg_q, + True, + D_lg, + None, + tex.DType.kBFloat16, + None, + tex.DType.kBFloat16, + False, + None, + False, + ws, + ws_size, + False, + False, + ) + + tflops, avg_ms = _time_with_cuda_events(_run, num_iters, flops, leading_fn=_lead) + del A_lg_q, B_lg_q, D_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="NVFP4") + except Exception as e: + print(f"Warning: FP4 prequantized benchmark failed: {e}") + return None + + +# --------------------------------------------------------------------------- +# Shape helpers +# --------------------------------------------------------------------------- +def get_default_shapes() -> list[tuple[int, int, int]]: + """Default set of square matrix shapes for benchmarking.""" + return [ + (256, 256, 256), + (512, 512, 512), + (768, 768, 768), + (1024, 1024, 1024), + (1536, 1536, 1536), + (2048, 2048, 2048), + (3072, 3072, 3072), + (4096, 4096, 4096), + (6144, 6144, 6144), + (8192, 8192, 8192), + (16384, 16384, 16384), + ] + + +def parse_shapes_arg(shapes_arg: str) -> list[tuple[int, int, int]]: + """Parse ``--shapes`` into a list of (M, K, N) tuples. + + Accepts either square sizes (``1024,2048,4096``) or explicit + triplets (``8192x5120x10240,8192x10240x5120``), or a mix. + + Raises: + ValueError: On malformed input. + """ + items = [s.strip() for s in shapes_arg.split(",") if s.strip()] + if not items: + raise ValueError("Empty --shapes argument.") + + shapes: list[tuple[int, int, int]] = [] + for item in items: + if "x" in item: + parts = [p.strip() for p in item.lower().split("x")] + if len(parts) != 3: + raise ValueError(f"Invalid shape '{item}'. Expected 'MxKxN'.") + shapes.append((int(parts[0]), int(parts[1]), int(parts[2]))) + else: + size = int(item) + shapes.append((size, size, size)) + return shapes + + +def compute_gemm_shapes( + config: ModelConfig, +) -> tuple[ + list[tuple[str, int, int, int]], + list[tuple[str, int, int, int]], + list[tuple[str, int, int, int]], +]: + """Derive Fprop, Dgrad, and Wgrad GEMM shapes from a transformer model config. + + For forward Y = X @ W with shape (M, K, N): + - Dgrad: dX = dY @ Wᵀ → (M, N, K) (K and N swap) + - Wgrad: dW = Xᵀ @ dY → (K, M, N) (M moves to contraction axis) + + Returns: + (fprop_shapes, dgrad_shapes, wgrad_shapes) where each is a list of + (label, M, K, N) tuples. + """ + H = config.hidden_size + I = config.intermediate_size + M = config.micro_batch_size * config.sequence_length + + if H % config.num_attention_heads != 0: + raise ValueError( + f"hidden_size ({H}) must be divisible by " + f"num_attention_heads ({config.num_attention_heads})" + ) + + N_qkv = 3 * H + + fprop_shapes = [ + ("QKV Proj", M, H, N_qkv), + ("Attn Out", M, H, H), + ("MLP Up", M, H, I), + ("MLP Down", M, I, H), + ] + + dgrad_shapes = [ + ("QKV Proj (Dgrad)", M, N_qkv, H), + ("Attn Out (Dgrad)", M, H, H), + ("MLP Up (Dgrad)", M, I, H), + ("MLP Down (Dgrad)", M, H, I), + ] + + wgrad_shapes = [ + ("QKV Proj (Wgrad)", H, M, N_qkv), + ("Attn Out (Wgrad)", H, M, H), + ("MLP Up (Wgrad)", H, M, I), + ("MLP Down (Wgrad)", I, M, H), + ] + + return fprop_shapes, dgrad_shapes, wgrad_shapes + + +# --------------------------------------------------------------------------- +# GPU warmup +# --------------------------------------------------------------------------- +def warmup_gpu(duration_seconds: float = 5.0) -> None: + """Run sustained matmuls to stabilize GPU clocks before benchmarking.""" + print(f"Warming up GPU for {duration_seconds:.1f} seconds...") + device = torch.device("cuda") + A = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + B = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + + torch.cuda.synchronize() + t0 = time.time() + while time.time() - t0 < duration_seconds: + for _ in range(10): + torch.matmul(A, B) + torch.cuda.synchronize() + + del A, B + torch.cuda.empty_cache() + print("GPU warmup complete.\n") + + +# --------------------------------------------------------------------------- +# Main orchestrator +# --------------------------------------------------------------------------- +def run_benchmarks( + shapes: list[tuple[int, int, int]], + num_warmup: int = 10, + num_iters: int = 100, + include_fp8_current: bool = True, + include_fp8_delayed: bool = True, + include_fp8: bool = True, + include_fp8_block: bool = True, + include_fp4: bool = True, + gpu_warmup_seconds: float = 5.0, + pre_quantize: bool = False, + timing: str = "cuda-events", + profile_shape: Optional[int] = None, +) -> dict[str, list[float]]: + """Run GEMM benchmarks for every shape and enabled precision. + + Returns: + Dict mapping precision name to a list of TFLOPS values, one per shape. + """ + results: dict[str, list[float]] = { + "BF16": [], + "FP8Current": [], + "FP8Delayed": [], + "FP8Block": [], + "MXFP8": [], + "NVFP4": [], + } + time_results: dict[str, list[float]] = { + "BF16": [], + "FP8Current": [], + "FP8Delayed": [], + "FP8Block": [], + "MXFP8": [], + "NVFP4": [], + } + + has_blackwell = is_blackwell_available() + run_fp8_current = include_fp8_current and TE_AVAILABLE + # DelayedScaling has no pre-quantized variant: the recipe differs from CurrentScaling + # only in how the scaling factor is computed each step (via amax history), which is + # exactly the work pre-quantized mode skips. Enabling it here would silently fall back + # to the autocast path and plot a misleading bar, so omit it when pre-quantizing. + run_fp8_delayed = include_fp8_delayed and TE_AVAILABLE and not pre_quantize + run_fp8 = include_fp8 and TE_AVAILABLE + run_fp8_block = include_fp8_block and TE_AVAILABLE + run_fp4 = include_fp4 and TE_AVAILABLE and has_blackwell + + gpu_name = torch.cuda.get_device_name(0) + timing_label = ( + "torch.profiler (CUPTI kernel timestamps)" if timing == "profiler" else "CUDA events" + ) + + print(f"\nGEMM Benchmark on {gpu_name}") + print(f"Timing method: {timing_label}") + print(f"Warmup iterations: {num_warmup}, Timed iterations: {num_iters}") + if pre_quantize: + print("Mode: Pre-quantized inputs (raw kernel throughput)") + else: + print("Mode: Autocast (includes quantization overhead)") + if not has_blackwell and include_fp4: + print("Note: NVFP4 requires Blackwell (SM100+), skipping FP4 benchmarks") + + if profile_shape is not None: + shapes = [(profile_shape, profile_shape, profile_shape)] + print(f"\n*** PROFILING MODE: shape {profile_shape}x{profile_shape}x{profile_shape} ***") + print( + "*** Run with: nsys profile --capture-range=cudaProfilerApi python